{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "OSS-Socratic-Models-Robot-Pick-Place-Jun-16.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "u2E0aERL8YpW",
        "o29RKDf680GZ",
        "d2UdCCcpwbl0",
        "6YErUBSo3RT2",
        "Uyc6-l8IgQPS",
        "K9K0Bxh2Gj8p",
        "B7JYho2gHNES",
        "IIxTA71w3ZA0",
        "4rqB5LrY3dXZ",
        "XOb79YDd3ijm",
        "GO2qUyPblbdp"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2021 Google LLC. SPDX-License-Identifier: Apache-2.0"
      ],
      "metadata": {
        "id": "u2E0aERL8YpW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Copyright 2021 Google LLC. SPDX-License-Identifier: Apache-2.0\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
      ],
      "metadata": {
        "id": "xgKBf6Ot76Y5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Socratic Models:** Robot Pick & Place\n",
        "\n",
        "[Socratic Models](https://socraticmodels.github.io) (SMs) is a framework that composes multiple pre-existing \"foundation models\" (e.g., large language models, visual language models, audio-language models) via prompt engineering to provide results for new multimodal tasks, without any model finetuning.\n",
        "\n",
        "This colab runs an example of SMs for robot pick and place (one of the tasks demonstrated in the work):\n",
        "\n",
        "<img src=\"https://socraticmodels.github.io/images/robot_demo_teaser.png\" height=\"320px\">\n",
        "\n",
        "### **Quick Start:**\n",
        "\n",
        "**Step 1.** Register for an [OpenAI API key](https://openai.com/blog/openai-api/) to use GPT-3 (there's a free trial) and enter it below\n",
        "\n",
        "**Step 2.** Menu > Change runtime type > Hardware accelerator > \"GPU\"\n",
        "\n",
        "**Step 3.** Menu > Runtime > Run all"
      ],
      "metadata": {
        "id": "o29RKDf680GZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "openai_api_key = \"your-api-key\""
      ],
      "metadata": {
        "id": "PZdNE66A-CrF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Setup**\n",
        "\n",
        "This does a few things:\n",
        "* Installs JAX, FLAX (learning), PyBullet (simulator), MoviePy (videos).\n",
        "* Imports Python packages and sets OpenAI API key.\n",
        "* Downloads PyBullet assets and ViLD pretrained model weights.\n",
        "* Loads TensorBoard extension."
      ],
      "metadata": {
        "id": "d2UdCCcpwbl0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown \n",
        "\n",
        "!pip install ftfy regex tqdm fvcore imageio imageio-ffmpeg\n",
        "!pip install git+https://github.com/openai/CLIP.git\n",
        "!pip install -U --no-cache-dir gdown --pre\n",
        "!pip install pybullet moviepy\n",
        "!pip install flax\n",
        "!pip install openai\n",
        "!pip install easydict\n",
        "# !pip install tensorflow==2.7.0  # If error: UNIMPLEMENTED: DNN library is not found.\n",
        "\n",
        "import collections\n",
        "import datetime\n",
        "import os\n",
        "import random\n",
        "import threading\n",
        "import time\n",
        "\n",
        "import cv2  # Used by ViLD.\n",
        "import clip\n",
        "from easydict import EasyDict\n",
        "import flax\n",
        "from flax import linen as nn\n",
        "from flax.training import checkpoints\n",
        "from flax.metrics import tensorboard\n",
        "import imageio\n",
        "import IPython\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "from moviepy.editor import ImageSequenceClip\n",
        "import numpy as np\n",
        "import openai\n",
        "import optax\n",
        "import pickle\n",
        "from PIL import Image\n",
        "import pybullet\n",
        "import pybullet_data\n",
        "import tensorflow.compat.v1 as tf\n",
        "import torch\n",
        "from tqdm import tqdm\n",
        "\n",
        "#Download PyBullet assets.\n",
        "if not os.path.exists('ur5e/ur5e.urdf'):\n",
        "  !gdown --id 1Cc_fDSBL6QiDvNT4dpfAEbhbALSVoWcc\n",
        "  !gdown --id 1yOMEm-Zp_DL3nItG9RozPeJAmeOldekX\n",
        "  !gdown --id 1GsqNLhEl9dd4Mc3BM0dX3MibOI1FVWNM\n",
        "  !unzip ur5e.zip\n",
        "  !unzip robotiq_2f_85.zip\n",
        "  !unzip bowl.zip\n",
        "\n",
        "# ViLD pretrained model weights.\n",
        "!gsutil cp -r gs://cloud-tpu-checkpoints/detection/projects/vild/colab/image_path_v2 ./\n",
        "\n",
        "%load_ext tensorboard\n",
        "\n",
        "# Set OpenAI API key.\n",
        "openai.api_key = openai_api_key\n",
        "\n",
        "# Show useful GPU info.\n",
        "!nvidia-smi\n",
        "\n",
        "# Show if JAX is using GPU.\n",
        "from jax.lib import xla_bridge\n",
        "print(xla_bridge.get_backend().platform)"
      ],
      "metadata": {
        "id": "NBSzsHw89qMp",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Environment**\n",
        "Define PyBullet-based environment with a UR5e and Robotiq 2F85 gripper.\n"
      ],
      "metadata": {
        "id": "6YErUBSo3RT2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown **Global constants:** pick and place objects, colors, workspace bounds.\n",
        "\n",
        "PICK_TARGETS = {\n",
        "  'blue block': None,\n",
        "  'red block': None,\n",
        "  'green block': None,\n",
        "  'orange block': None,\n",
        "  'yellow block': None,\n",
        "  'purple block': None,\n",
        "  'pink block': None,\n",
        "  'cyan block': None,\n",
        "  'brown block': None,\n",
        "  'gray block': None,\n",
        "}\n",
        "\n",
        "COLORS = {\n",
        "  'blue':   (78/255,  121/255, 167/255, 255/255),\n",
        "  'red':    (255/255,  87/255,  89/255, 255/255),\n",
        "  'green':  (89/255,  169/255,  79/255, 255/255),\n",
        "  'orange': (242/255, 142/255,  43/255, 255/255),\n",
        "  'yellow': (237/255, 201/255,  72/255, 255/255),\n",
        "  'purple': (176/255, 122/255, 161/255, 255/255),\n",
        "  'pink':   (255/255, 157/255, 167/255, 255/255),\n",
        "  'cyan':   (118/255, 183/255, 178/255, 255/255),\n",
        "  'brown':  (156/255, 117/255,  95/255, 255/255),\n",
        "  'gray':   (186/255, 176/255, 172/255, 255/255),\n",
        "}\n",
        "\n",
        "PLACE_TARGETS = {\n",
        "  'blue block': None,\n",
        "  'red block': None,\n",
        "  'green block': None,\n",
        "  'orange block': None,\n",
        "  'yellow block': None,\n",
        "  'purple block': None,\n",
        "  'pink block': None,\n",
        "  'cyan block': None,\n",
        "  'brown block': None,\n",
        "  'gray block': None,\n",
        "\n",
        "  'blue bowl': None,\n",
        "  'red bowl': None,\n",
        "  'green bowl': None,\n",
        "  'orange bowl': None,\n",
        "  'yellow bowl': None,\n",
        "  'purple bowl': None,\n",
        "  'pink bowl': None,\n",
        "  'cyan bowl': None,\n",
        "  'brown bowl': None,\n",
        "  'gray bowl': None,\n",
        "\n",
        "  'top left corner':     (-0.3 + 0.05, -0.2 - 0.05, 0),\n",
        "  'top side':            (0,           -0.2 - 0.05, 0),\n",
        "  'top right corner':    (0.3 - 0.05,  -0.2 - 0.05, 0),\n",
        "  'left side':           (-0.3 + 0.05, -0.5,        0),\n",
        "  'middle':              (0,           -0.5,        0),\n",
        "  'right side':          (0.3 - 0.05,  -0.5,        0),\n",
        "  'bottom left corner':  (-0.3 + 0.05, -0.8 + 0.05, 0),\n",
        "  'bottom side':         (0,           -0.8 + 0.05, 0),\n",
        "  'bottom right corner': (0.3 - 0.05,  -0.8 + 0.05, 0),\n",
        "}\n",
        "\n",
        "PIXEL_SIZE = 0.00267857\n",
        "BOUNDS = np.float32([[-0.3, 0.3], [-0.8, -0.2], [0, 0.15]])  # (X, Y, Z)"
      ],
      "metadata": {
        "id": "D6gTAdeX39yd",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown **Gripper class:** adds a gripper to the robot and runs a parallel thread to simulate single-actuator behavior.\n",
        "\n",
        "class Robotiq2F85:\n",
        "  \"\"\"Gripper handling for Robotiq 2F85.\"\"\"\n",
        "\n",
        "  def __init__(self, robot, tool):\n",
        "    self.robot = robot\n",
        "    self.tool = tool\n",
        "    pos = [0.1339999999999999, -0.49199999999872496, 0.5]\n",
        "    rot = pybullet.getQuaternionFromEuler([np.pi, 0, np.pi])\n",
        "    urdf = 'robotiq_2f_85/robotiq_2f_85.urdf'\n",
        "    self.body = pybullet.loadURDF(urdf, pos, rot)\n",
        "    self.n_joints = pybullet.getNumJoints(self.body)\n",
        "    self.activated = False\n",
        "\n",
        "    # Connect gripper base to robot tool.\n",
        "    pybullet.createConstraint(self.robot, tool, self.body, 0, jointType=pybullet.JOINT_FIXED, jointAxis=[0, 0, 0], parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, -0.07], childFrameOrientation=pybullet.getQuaternionFromEuler([0, 0, np.pi / 2]))\n",
        "\n",
        "    # Set friction coefficients for gripper fingers.\n",
        "    for i in range(pybullet.getNumJoints(self.body)):\n",
        "      pybullet.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)\n",
        "\n",
        "    # Start thread to handle additional gripper constraints.\n",
        "    self.motor_joint = 1\n",
        "    self.running = True\n",
        "    self.constraints_thread = threading.Thread(target=self.step)\n",
        "    self.constraints_thread.daemon = True\n",
        "    self.constraints_thread.start()\n",
        "\n",
        "  # Control joint positions by enforcing hard contraints on gripper behavior.\n",
        "  # Set one joint as the open/close motor joint (other joints should mimic).\n",
        "  def step(self):\n",
        "    while self.running:\n",
        "      try:\n",
        "        currj = [pybullet.getJointState(self.body, i)[0] for i in range(self.n_joints)]\n",
        "        indj = [6, 3, 8, 5, 10]\n",
        "        targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]\n",
        "        pybullet.setJointMotorControlArray(self.body, indj, pybullet.POSITION_CONTROL, targj, positionGains=np.ones(5))\n",
        "      except:\n",
        "        return\n",
        "      time.sleep(0.001)\n",
        "\n",
        "  # Close gripper fingers.\n",
        "  def activate(self):\n",
        "    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=1, force=10)\n",
        "    self.activated = True\n",
        "\n",
        "  # Open gripper fingers.\n",
        "  def release(self):\n",
        "    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=-1, force=10)\n",
        "    self.activated = False\n",
        "\n",
        "  # If activated and object in gripper: check object contact.\n",
        "  # If activated and nothing in gripper: check gripper contact.\n",
        "  # If released: check proximity to surface (disabled).\n",
        "  def detect_contact(self):\n",
        "    obj, _, ray_frac = self.check_proximity()\n",
        "    if self.activated:\n",
        "      empty = self.grasp_width() < 0.01\n",
        "      cbody = self.body if empty else obj\n",
        "      if obj == self.body or obj == 0:\n",
        "        return False\n",
        "      return self.external_contact(cbody)\n",
        "  #   else:\n",
        "  #     return ray_frac < 0.14 or self.external_contact()\n",
        "\n",
        "  # Return if body is in contact with something other than gripper\n",
        "  def external_contact(self, body=None):\n",
        "    if body is None:\n",
        "      body = self.body\n",
        "    pts = pybullet.getContactPoints(bodyA=body)\n",
        "    pts = [pt for pt in pts if pt[2] != self.body]\n",
        "    return len(pts) > 0  # pylint: disable=g-explicit-length-test\n",
        "\n",
        "  def check_grasp(self):\n",
        "    while self.moving():\n",
        "      time.sleep(0.001)\n",
        "    success = self.grasp_width() > 0.01\n",
        "    return success\n",
        "\n",
        "  def grasp_width(self):\n",
        "    lpad = np.array(pybullet.getLinkState(self.body, 4)[0])\n",
        "    rpad = np.array(pybullet.getLinkState(self.body, 9)[0])\n",
        "    dist = np.linalg.norm(lpad - rpad) - 0.047813\n",
        "    return dist\n",
        "\n",
        "  def check_proximity(self):\n",
        "    ee_pos = np.array(pybullet.getLinkState(self.robot, self.tool)[0])\n",
        "    tool_pos = np.array(pybullet.getLinkState(self.body, 0)[0])\n",
        "    vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))\n",
        "    ee_targ = ee_pos + vec\n",
        "    ray_data = pybullet.rayTest(ee_pos, ee_targ)[0]\n",
        "    obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]\n",
        "    return obj, link, ray_frac"
      ],
      "metadata": {
        "id": "0bj91lBq6WeK",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown **Gym-style environment class:** this initializes a robot overlooking a workspace with objects.\n",
        "\n",
        "class PickPlaceEnv():\n",
        "\n",
        "  def __init__(self):\n",
        "    self.dt = 1/480\n",
        "    self.sim_step = 0\n",
        "\n",
        "    # Configure and start PyBullet.\n",
        "    # python3 -m pybullet_utils.runServer\n",
        "    # pybullet.connect(pybullet.SHARED_MEMORY)  # pybullet.GUI for local GUI.\n",
        "    pybullet.connect(pybullet.DIRECT)  # pybullet.GUI for local GUI.\n",
        "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)\n",
        "    pybullet.setPhysicsEngineParameter(enableFileCaching=0)\n",
        "    assets_path = os.path.dirname(os.path.abspath(\"\"))\n",
        "    pybullet.setAdditionalSearchPath(assets_path)\n",
        "    pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())\n",
        "    pybullet.setTimeStep(self.dt)\n",
        "\n",
        "    self.home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0)  # Joint angles: (J0, J1, J2, J3, J4, J5).\n",
        "    self.home_ee_euler = (np.pi, 0, np.pi)  # (RX, RY, RZ) rotation in Euler angles.\n",
        "    self.ee_link_id = 9  # Link ID of UR5 end effector.\n",
        "    self.tip_link_id = 10  # Link ID of gripper finger tips.\n",
        "    self.gripper = None\n",
        "\n",
        "  def reset(self, config):\n",
        "    pybullet.resetSimulation(pybullet.RESET_USE_DEFORMABLE_WORLD)\n",
        "    pybullet.setGravity(0, 0, -9.8)\n",
        "    self.cache_video = []\n",
        "\n",
        "    # Temporarily disable rendering to load URDFs faster.\n",
        "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)\n",
        "\n",
        "    # Add robot.\n",
        "    pybullet.loadURDF(\"plane.urdf\", [0, 0, -0.001])\n",
        "    self.robot_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, 0], flags=pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)\n",
        "    self.ghost_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, -10])  # For forward kinematics.\n",
        "    self.joint_ids = [pybullet.getJointInfo(self.robot_id, i) for i in range(pybullet.getNumJoints(self.robot_id))]\n",
        "    self.joint_ids = [j[0] for j in self.joint_ids if j[2] == pybullet.JOINT_REVOLUTE]\n",
        "\n",
        "    # Move robot to home configuration.\n",
        "    for i in range(len(self.joint_ids)):\n",
        "      pybullet.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])\n",
        "\n",
        "    # Add gripper.\n",
        "    if self.gripper is not None:\n",
        "      while self.gripper.constraints_thread.is_alive():\n",
        "        self.constraints_thread_active = False\n",
        "    self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id)\n",
        "    self.gripper.release()\n",
        "\n",
        "    # Add workspace.\n",
        "    plane_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n",
        "    plane_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n",
        "    plane_id = pybullet.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])\n",
        "    pybullet.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])\n",
        "\n",
        "    # Load objects according to config.\n",
        "    self.config = config\n",
        "    self.obj_name_to_id = {}\n",
        "    obj_names = list(self.config['pick']) + list(self.config['place'])\n",
        "    obj_xyz = np.zeros((0, 3))\n",
        "    for obj_name in obj_names:\n",
        "      if ('block' in obj_name) or ('bowl' in obj_name):\n",
        "\n",
        "        # Get random position 15cm+ from other objects.\n",
        "        while True:\n",
        "          rand_x = np.random.uniform(BOUNDS[0, 0] + 0.1, BOUNDS[0, 1] - 0.1)\n",
        "          rand_y = np.random.uniform(BOUNDS[1, 0] + 0.1, BOUNDS[1, 1] - 0.1)\n",
        "          rand_xyz = np.float32([rand_x, rand_y, 0.03]).reshape(1, 3)\n",
        "          if len(obj_xyz) == 0:\n",
        "            obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n",
        "            break\n",
        "          else:\n",
        "            nn_dist = np.min(np.linalg.norm(obj_xyz - rand_xyz, axis=1)).squeeze()\n",
        "            if nn_dist > 0.15:\n",
        "              obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n",
        "              break\n",
        "        \n",
        "        object_color = COLORS[obj_name.split(' ')[0]]\n",
        "        object_type = obj_name.split(' ')[1]\n",
        "        object_position = rand_xyz.squeeze()\n",
        "        if object_type == 'block':\n",
        "          object_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n",
        "          object_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n",
        "          object_id = pybullet.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)\n",
        "        elif object_type == 'bowl':\n",
        "          object_position[2] = 0\n",
        "          object_id = pybullet.loadURDF(\"bowl/bowl.urdf\", object_position, useFixedBase=1)\n",
        "        pybullet.changeVisualShape(object_id, -1, rgbaColor=object_color)\n",
        "        self.obj_name_to_id[obj_name] = object_id\n",
        "\n",
        "    # Re-enable rendering.\n",
        "    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)\n",
        "\n",
        "    for _ in range(200):\n",
        "      pybullet.stepSimulation()\n",
        "    print('Environment reset: done.')\n",
        "    return self.get_observation()\n",
        "\n",
        "  def servoj(self, joints):\n",
        "    \"\"\"Move to target joint positions with position control.\"\"\"\n",
        "    pybullet.setJointMotorControlArray(\n",
        "      bodyIndex=self.robot_id,\n",
        "      jointIndices=self.joint_ids,\n",
        "      controlMode=pybullet.POSITION_CONTROL,\n",
        "      targetPositions=joints,\n",
        "      positionGains=[0.01]*6)\n",
        "  \n",
        "  def movep(self, position):\n",
        "    \"\"\"Move to target end effector position.\"\"\"\n",
        "    joints = pybullet.calculateInverseKinematics(\n",
        "        bodyUniqueId=self.robot_id,\n",
        "        endEffectorLinkIndex=self.tip_link_id,\n",
        "        targetPosition=position,\n",
        "        targetOrientation=pybullet.getQuaternionFromEuler(self.home_ee_euler),\n",
        "        maxNumIterations=100)\n",
        "    self.servoj(joints)\n",
        "\n",
        "  def step(self, action=None):\n",
        "    \"\"\"Do pick and place motion primitive.\"\"\"\n",
        "    pick_xyz, place_xyz = action['pick'].copy(), action['place'].copy()\n",
        "\n",
        "    # Set fixed primitive z-heights.\n",
        "    hover_xyz = pick_xyz.copy() + np.float32([0, 0, 0.2])\n",
        "    pick_xyz[2] -= 0.02\n",
        "    pick_xyz[2] = max(pick_xyz[2], 0.03)\n",
        "    place_xyz[2] = 0.15\n",
        "\n",
        "    # Move to object.\n",
        "    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(hover_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "    while np.linalg.norm(pick_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(pick_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "\n",
        "    # Pick up object.\n",
        "    self.gripper.activate()\n",
        "    for _ in range(240):\n",
        "      self.step_sim_and_render()\n",
        "    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(hover_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "    \n",
        "    # Move to place location.\n",
        "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(place_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "\n",
        "    # Place down object.\n",
        "    while (not self.gripper.detect_contact()) and (place_xyz[2] > 0.03):\n",
        "      place_xyz[2] -= 0.001\n",
        "      self.movep(place_xyz)\n",
        "      for _ in range(3):\n",
        "        self.step_sim_and_render()\n",
        "    self.gripper.release()\n",
        "    for _ in range(240):\n",
        "      self.step_sim_and_render()\n",
        "    place_xyz[2] = 0.2\n",
        "    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(place_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "    place_xyz = np.float32([0, -0.5, 0.2])\n",
        "    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n",
        "      self.movep(place_xyz)\n",
        "      self.step_sim_and_render()\n",
        "      ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n",
        "\n",
        "    observation = self.get_observation()\n",
        "    reward = self.get_reward()\n",
        "    done = False\n",
        "    info = {}\n",
        "    return observation, reward, done, info\n",
        "\n",
        "  def step_sim_and_render(self):\n",
        "    pybullet.stepSimulation()\n",
        "    self.sim_step += 1\n",
        "\n",
        "    # Render current image at 8 FPS.\n",
        "    if self.sim_step % (1 / (8 * self.dt)) == 0:\n",
        "      self.cache_video.append(self.get_camera_image())\n",
        "\n",
        "  def get_camera_image(self):\n",
        "    image_size = (240, 240)\n",
        "    intrinsics = (120., 0, 120., 0, 120., 120., 0, 0, 1)\n",
        "    color, _, _, _, _ = env.render_image(image_size, intrinsics)\n",
        "    return color\n",
        "\n",
        "  def set_alpha_transparency(self, alpha: float) -> None:\n",
        "    for id in range(20):\n",
        "      visual_shape_data = pybullet.getVisualShapeData(id)\n",
        "      for i in range(len(visual_shape_data)):\n",
        "        object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]\n",
        "        rgba_color = list(rgba_color[0:3]) +  [alpha]\n",
        "        pybullet.changeVisualShape(self.robot_id, linkIndex=i, rgbaColor=rgba_color)      \n",
        "        pybullet.changeVisualShape(self.gripper.body, linkIndex=i, rgbaColor=rgba_color)\n",
        "\n",
        "  def get_camera_image_top(self, \n",
        "                           image_size=(240, 240), \n",
        "                           intrinsics=(2000., 0, 2000., 0, 2000., 2000., 0, 0, 1),\n",
        "                           position=(0, -0.5, 5),\n",
        "                           orientation=(0, np.pi, -np.pi / 2),\n",
        "                           zrange=(0.01, 1.),\n",
        "                           set_alpha=True):\n",
        "    set_alpha and self.set_alpha_transparency(0)\n",
        "    color, _, _, _, _ = env.render_image_top(image_size, \n",
        "                                             intrinsics,\n",
        "                                             position,\n",
        "                                             orientation,\n",
        "                                             zrange)\n",
        "    set_alpha and self.set_alpha_transparency(1)\n",
        "    return color\n",
        "\n",
        "  def render_image_top(self, \n",
        "                       image_size=(240, 240), \n",
        "                       intrinsics=(2000., 0, 2000., 0, 2000., 2000., 0, 0, 1),\n",
        "                       position=(0, -0.5, 5),\n",
        "                       orientation=(0, np.pi, -np.pi / 2),\n",
        "                       zrange=(0.01, 1.)):\n",
        "\n",
        "    # Camera parameters.\n",
        "    orientation = pybullet.getQuaternionFromEuler(orientation)\n",
        "    noise=True\n",
        "\n",
        "    # OpenGL camera settings.\n",
        "    lookdir = np.float32([0, 0, 1]).reshape(3, 1)\n",
        "    updir = np.float32([0, -1, 0]).reshape(3, 1)\n",
        "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
        "    rotm = np.float32(rotation).reshape(3, 3)\n",
        "    lookdir = (rotm @ lookdir).reshape(-1)\n",
        "    updir = (rotm @ updir).reshape(-1)\n",
        "    lookat = position + lookdir\n",
        "    focal_len = intrinsics[0]\n",
        "    znear, zfar = (0.01, 10.)\n",
        "    viewm = pybullet.computeViewMatrix(position, lookat, updir)\n",
        "    fovh = (image_size[0] / 2) / focal_len\n",
        "    fovh = 180 * np.arctan(fovh) * 2 / np.pi\n",
        "\n",
        "    # Notes: 1) FOV is vertical FOV 2) aspect must be float\n",
        "    aspect_ratio = image_size[1] / image_size[0]\n",
        "    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)\n",
        "\n",
        "    # Render with OpenGL camera settings.\n",
        "    _, _, color, depth, segm = pybullet.getCameraImage(\n",
        "        width=image_size[1],\n",
        "        height=image_size[0],\n",
        "        viewMatrix=viewm,\n",
        "        projectionMatrix=projm,\n",
        "        shadow=1,\n",
        "        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,\n",
        "        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)\n",
        "\n",
        "    # Get color image.\n",
        "    color_image_size = (image_size[0], image_size[1], 4)\n",
        "    color = np.array(color, dtype=np.uint8).reshape(color_image_size)\n",
        "    color = color[:, :, :3]  # remove alpha channel\n",
        "    if noise:\n",
        "      color = np.int32(color)\n",
        "      color += np.int32(np.random.normal(0, 3, color.shape))\n",
        "      color = np.uint8(np.clip(color, 0, 255))\n",
        "\n",
        "    # Get depth image.\n",
        "    depth_image_size = (image_size[0], image_size[1])\n",
        "    zbuffer = np.float32(depth).reshape(depth_image_size)\n",
        "    depth = (zfar + znear - (2 * zbuffer - 1) * (zfar - znear))\n",
        "    depth = (2 * znear * zfar) / depth\n",
        "    if noise:\n",
        "      depth += np.random.normal(0, 0.003, depth.shape)\n",
        "\n",
        "    intrinsics = np.float32(intrinsics).reshape(3, 3)\n",
        "    return color, depth, position, orientation, intrinsics\n",
        "\n",
        "  def get_reward(self):\n",
        "    return 0  # TODO: check did the robot follow text instructions?\n",
        "\n",
        "  def get_observation(self):\n",
        "    observation = {}\n",
        "\n",
        "    # Render current image.\n",
        "    color, depth, position, orientation, intrinsics = self.render_image()\n",
        "\n",
        "    # Get heightmaps and colormaps.\n",
        "    points = self.get_pointcloud(depth, intrinsics)\n",
        "    position = np.float32(position).reshape(3, 1)\n",
        "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
        "    rotation = np.float32(rotation).reshape(3, 3)\n",
        "    transform = np.eye(4)\n",
        "    transform[:3, :] = np.hstack((rotation, position))\n",
        "    points = self.transform_pointcloud(points, transform)\n",
        "    heightmap, colormap, xyzmap = self.get_heightmap(points, color, BOUNDS, PIXEL_SIZE)\n",
        "\n",
        "    observation[\"image\"] = colormap\n",
        "    observation[\"xyzmap\"] = xyzmap\n",
        "    return observation\n",
        "\n",
        "  def render_image(self, image_size=(720, 720), intrinsics=(360., 0, 360., 0, 360., 360., 0, 0, 1)):\n",
        "\n",
        "    # Camera parameters.\n",
        "    position = (0, -0.85, 0.4)\n",
        "    orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)\n",
        "    orientation = pybullet.getQuaternionFromEuler(orientation)\n",
        "    zrange = (0.01, 10.)\n",
        "    noise=True\n",
        "\n",
        "    # OpenGL camera settings.\n",
        "    lookdir = np.float32([0, 0, 1]).reshape(3, 1)\n",
        "    updir = np.float32([0, -1, 0]).reshape(3, 1)\n",
        "    rotation = pybullet.getMatrixFromQuaternion(orientation)\n",
        "    rotm = np.float32(rotation).reshape(3, 3)\n",
        "    lookdir = (rotm @ lookdir).reshape(-1)\n",
        "    updir = (rotm @ updir).reshape(-1)\n",
        "    lookat = position + lookdir\n",
        "    focal_len = intrinsics[0]\n",
        "    znear, zfar = (0.01, 10.)\n",
        "    viewm = pybullet.computeViewMatrix(position, lookat, updir)\n",
        "    fovh = (image_size[0] / 2) / focal_len\n",
        "    fovh = 180 * np.arctan(fovh) * 2 / np.pi\n",
        "\n",
        "    # Notes: 1) FOV is vertical FOV 2) aspect must be float\n",
        "    aspect_ratio = image_size[1] / image_size[0]\n",
        "    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)\n",
        "\n",
        "    # Render with OpenGL camera settings.\n",
        "    _, _, color, depth, segm = pybullet.getCameraImage(\n",
        "        width=image_size[1],\n",
        "        height=image_size[0],\n",
        "        viewMatrix=viewm,\n",
        "        projectionMatrix=projm,\n",
        "        shadow=1,\n",
        "        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,\n",
        "        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)\n",
        "\n",
        "    # Get color image.\n",
        "    color_image_size = (image_size[0], image_size[1], 4)\n",
        "    color = np.array(color, dtype=np.uint8).reshape(color_image_size)\n",
        "    color = color[:, :, :3]  # remove alpha channel\n",
        "    if noise:\n",
        "      color = np.int32(color)\n",
        "      color += np.int32(np.random.normal(0, 3, color.shape))\n",
        "      color = np.uint8(np.clip(color, 0, 255))\n",
        "\n",
        "    # Get depth image.\n",
        "    depth_image_size = (image_size[0], image_size[1])\n",
        "    zbuffer = np.float32(depth).reshape(depth_image_size)\n",
        "    depth = (zfar + znear - (2 * zbuffer - 1) * (zfar - znear))\n",
        "    depth = (2 * znear * zfar) / depth\n",
        "    if noise:\n",
        "      depth += np.random.normal(0, 0.003, depth.shape)\n",
        "\n",
        "    intrinsics = np.float32(intrinsics).reshape(3, 3)\n",
        "    return color, depth, position, orientation, intrinsics\n",
        "\n",
        "  def get_pointcloud(self, depth, intrinsics):\n",
        "    \"\"\"Get 3D pointcloud from perspective depth image.\n",
        "    Args:\n",
        "      depth: HxW float array of perspective depth in meters.\n",
        "      intrinsics: 3x3 float array of camera intrinsics matrix.\n",
        "    Returns:\n",
        "      points: HxWx3 float array of 3D points in camera coordinates.\n",
        "    \"\"\"\n",
        "    height, width = depth.shape\n",
        "    xlin = np.linspace(0, width - 1, width)\n",
        "    ylin = np.linspace(0, height - 1, height)\n",
        "    px, py = np.meshgrid(xlin, ylin)\n",
        "    px = (px - intrinsics[0, 2]) * (depth / intrinsics[0, 0])\n",
        "    py = (py - intrinsics[1, 2]) * (depth / intrinsics[1, 1])\n",
        "    points = np.float32([px, py, depth]).transpose(1, 2, 0)\n",
        "    return points\n",
        "\n",
        "  def transform_pointcloud(self, points, transform):\n",
        "    \"\"\"Apply rigid transformation to 3D pointcloud.\n",
        "    Args:\n",
        "      points: HxWx3 float array of 3D points in camera coordinates.\n",
        "      transform: 4x4 float array representing a rigid transformation matrix.\n",
        "    Returns:\n",
        "      points: HxWx3 float array of transformed 3D points.\n",
        "    \"\"\"\n",
        "    padding = ((0, 0), (0, 0), (0, 1))\n",
        "    homogen_points = np.pad(points.copy(), padding,\n",
        "                            'constant', constant_values=1)\n",
        "    for i in range(3):\n",
        "      points[Ellipsis, i] = np.sum(transform[i, :] * homogen_points, axis=-1)\n",
        "    return points\n",
        "\n",
        "  def get_heightmap(self, points, colors, bounds, pixel_size):\n",
        "    \"\"\"Get top-down (z-axis) orthographic heightmap image from 3D pointcloud.\n",
        "    Args:\n",
        "      points: HxWx3 float array of 3D points in world coordinates.\n",
        "      colors: HxWx3 uint8 array of values in range 0-255 aligned with points.\n",
        "      bounds: 3x2 float array of values (rows: X,Y,Z; columns: min,max) defining\n",
        "        region in 3D space to generate heightmap in world coordinates.\n",
        "      pixel_size: float defining size of each pixel in meters.\n",
        "    Returns:\n",
        "      heightmap: HxW float array of height (from lower z-bound) in meters.\n",
        "      colormap: HxWx3 uint8 array of backprojected color aligned with heightmap.\n",
        "      xyzmap: HxWx3 float array of XYZ points in world coordinates.\n",
        "    \"\"\"\n",
        "    width = int(np.round((bounds[0, 1] - bounds[0, 0]) / pixel_size))\n",
        "    height = int(np.round((bounds[1, 1] - bounds[1, 0]) / pixel_size))\n",
        "    heightmap = np.zeros((height, width), dtype=np.float32)\n",
        "    colormap = np.zeros((height, width, colors.shape[-1]), dtype=np.uint8)\n",
        "    xyzmap = np.zeros((height, width, 3), dtype=np.float32)\n",
        "\n",
        "    # Filter out 3D points that are outside of the predefined bounds.\n",
        "    ix = (points[Ellipsis, 0] >= bounds[0, 0]) & (points[Ellipsis, 0] < bounds[0, 1])\n",
        "    iy = (points[Ellipsis, 1] >= bounds[1, 0]) & (points[Ellipsis, 1] < bounds[1, 1])\n",
        "    iz = (points[Ellipsis, 2] >= bounds[2, 0]) & (points[Ellipsis, 2] < bounds[2, 1])\n",
        "    valid = ix & iy & iz\n",
        "    points = points[valid]\n",
        "    colors = colors[valid]\n",
        "\n",
        "    # Sort 3D points by z-value, which works with array assignment to simulate\n",
        "    # z-buffering for rendering the heightmap image.\n",
        "    iz = np.argsort(points[:, -1])\n",
        "    points, colors = points[iz], colors[iz]\n",
        "    px = np.int32(np.floor((points[:, 0] - bounds[0, 0]) / pixel_size))\n",
        "    py = np.int32(np.floor((points[:, 1] - bounds[1, 0]) / pixel_size))\n",
        "    px = np.clip(px, 0, width - 1)\n",
        "    py = np.clip(py, 0, height - 1)\n",
        "    heightmap[py, px] = points[:, 2] - bounds[2, 0]\n",
        "    for c in range(colors.shape[-1]):\n",
        "      colormap[py, px, c] = colors[:, c]\n",
        "      xyzmap[py, px, c] = points[:, c]\n",
        "    colormap = colormap[::-1, :, :]  # Flip up-down.\n",
        "    xv, yv = np.meshgrid(np.linspace(BOUNDS[0, 0], BOUNDS[0, 1], height),\n",
        "                         np.linspace(BOUNDS[1, 0], BOUNDS[1, 1], width))\n",
        "    xyzmap[:, :, 0] = xv\n",
        "    xyzmap[:, :, 1] = yv\n",
        "    xyzmap = xyzmap[::-1, :, :]  # Flip up-down.\n",
        "    heightmap = heightmap[::-1, :]  # Flip up-down.\n",
        "    return heightmap, colormap, xyzmap\n",
        "\n",
        "def xyz_to_pix(position):\n",
        "  \"\"\"Convert from 3D position to pixel location on heightmap.\"\"\"\n",
        "  u = int(np.round((BOUNDS[1, 1] - position[1]) / PIXEL_SIZE))\n",
        "  v = int(np.round((position[0] - BOUNDS[0, 0]) / PIXEL_SIZE))\n",
        "  return (u, v)"
      ],
      "metadata": {
        "id": "ELFr8iE28hVX",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Initialize an environment and render images.\n",
        "\n",
        "if 'env' in locals():\n",
        "  # Safely exit gripper threading before re-initializing environment.\n",
        "  env.gripper.running = False\n",
        "  while env.gripper.constraints_thread.isAlive():\n",
        "    time.sleep(0.01)\n",
        "env = PickPlaceEnv()\n",
        "\n",
        "# Define and reset environment.\n",
        "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
        "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
        "\n",
        "np.random.seed(42)\n",
        "obs = env.reset(config)\n",
        "\n",
        "plt.subplot(1, 2, 1)\n",
        "img = env.get_camera_image()\n",
        "plt.title('Perspective side-view')\n",
        "plt.imshow(img)\n",
        "plt.subplot(1, 2, 2)\n",
        "img = env.get_camera_image_top()\n",
        "img = np.flipud(img.transpose(1, 0, 2))\n",
        "plt.title('Orthographic top-view')\n",
        "plt.imshow(img)\n",
        "plt.show()\n",
        "\n",
        "# Note: orthographic cameras do not exist. But we can approximate them by\n",
        "# projecting a 3D point cloud from an RGB-D camera, then unprojecting that onto\n",
        "# an orthographic plane. Orthographic views are useful for spatial action maps.\n",
        "plt.title('Unprojected orthographic top-view')\n",
        "plt.imshow(obs['image'])\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "HclXxMiNFHyE",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Demo:** ViLD\n",
        "Run zero-shot open-vocabulary object detection with [ViLD](https://arxiv.org/abs/2104.13921) to generate a list of objects as a scene description for a large language model."
      ],
      "metadata": {
        "id": "Uyc6-l8IgQPS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define and reset environment.\n",
        "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
        "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
        "\n",
        "np.random.seed(42)\n",
        "obs = env.reset(config)\n",
        "img = env.get_camera_image_top()\n",
        "img = np.flipud(img.transpose(1, 0, 2))\n",
        "plt.title('ViLD Input Image')\n",
        "plt.imshow(img)\n",
        "plt.show()\n",
        "imageio.imwrite('tmp.jpg', img)"
      ],
      "metadata": {
        "id": "DBkMaGRng1Vd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Load CLIP model.\n",
        "\n",
        "torch.cuda.set_per_process_memory_fraction(0.9, None)\n",
        "clip_model, clip_preprocess = clip.load(\"ViT-B/32\")\n",
        "clip_model.cuda().eval()\n",
        "print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in clip_model.parameters()]):,}\")\n",
        "print(\"Input resolution:\", clip_model.visual.input_resolution)\n",
        "print(\"Context length:\", clip_model.context_length)\n",
        "print(\"Vocab size:\", clip_model.vocab_size)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ckfJ6LwjEiIQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Define ViLD hyperparameters.\n",
        "FLAGS = {\n",
        "    'prompt_engineering': True,\n",
        "    'this_is': True,\n",
        "    'temperature': 100.0,\n",
        "    'use_softmax': False,\n",
        "}\n",
        "FLAGS = EasyDict(FLAGS)\n",
        "\n",
        "\n",
        "# # Global matplotlib settings\n",
        "# SMALL_SIZE = 16#10\n",
        "# MEDIUM_SIZE = 18#12\n",
        "# BIGGER_SIZE = 20#14\n",
        "\n",
        "# plt.rc('font', size=MEDIUM_SIZE)         # controls default text sizes\n",
        "# plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title\n",
        "# plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels\n",
        "# plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
        "# plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
        "# plt.rc('legend', fontsize=MEDIUM_SIZE)   # legend fontsize\n",
        "# plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title\n",
        "\n",
        "\n",
        "# Parameters for drawing figure.\n",
        "display_input_size = (10, 10)\n",
        "overall_fig_size = (18, 24)\n",
        "\n",
        "line_thickness = 1\n",
        "fig_size_w = 35\n",
        "# fig_size_h = min(max(5, int(len(category_names) / 2.5) ), 10)\n",
        "mask_color =   'red'\n",
        "alpha = 0.5"
      ],
      "metadata": {
        "cellView": "form",
        "id": "l-qWw2WDh_dp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ViLD prompt engineering.\n",
        "\n",
        "def article(name):\n",
        "  return \"an\" if name[0] in \"aeiou\" else \"a\"\n",
        "\n",
        "def processed_name(name, rm_dot=False):\n",
        "  # _ for lvis\n",
        "  # / for obj365\n",
        "  res = name.replace(\"_\", \" \").replace(\"/\", \" or \").lower()\n",
        "  if rm_dot:\n",
        "    res = res.rstrip(\".\")\n",
        "  return res\n",
        "\n",
        "single_template = [\n",
        "    \"a photo of {article} {}.\"\n",
        "]\n",
        "\n",
        "# multiple_templates = [\n",
        "#     \"There is {article} {} in the scene.\",\n",
        "#     \"a painting of a {}.\",\n",
        "# ]\n",
        "\n",
        "multiple_templates = [\n",
        "    'There is {article} {} in the scene.',\n",
        "    'There is the {} in the scene.',\n",
        "    'a photo of {article} {} in the scene.',\n",
        "    'a photo of the {} in the scene.',\n",
        "    'a photo of one {} in the scene.',\n",
        "\n",
        "\n",
        "    'itap of {article} {}.',\n",
        "    'itap of my {}.',  # itap: I took a picture of\n",
        "    'itap of the {}.',\n",
        "    'a photo of {article} {}.',\n",
        "    'a photo of my {}.',\n",
        "    'a photo of the {}.',\n",
        "    'a photo of one {}.',\n",
        "    'a photo of many {}.',\n",
        "\n",
        "    'a good photo of {article} {}.',\n",
        "    'a good photo of the {}.',\n",
        "    'a bad photo of {article} {}.',\n",
        "    'a bad photo of the {}.',\n",
        "    'a photo of a nice {}.',\n",
        "    'a photo of the nice {}.',\n",
        "    'a photo of a cool {}.',\n",
        "    'a photo of the cool {}.',\n",
        "    'a photo of a weird {}.',\n",
        "    'a photo of the weird {}.',\n",
        "\n",
        "    'a photo of a small {}.',\n",
        "    'a photo of the small {}.',\n",
        "    'a photo of a large {}.',\n",
        "    'a photo of the large {}.',\n",
        "\n",
        "    'a photo of a clean {}.',\n",
        "    'a photo of the clean {}.',\n",
        "    'a photo of a dirty {}.',\n",
        "    'a photo of the dirty {}.',\n",
        "\n",
        "    'a bright photo of {article} {}.',\n",
        "    'a bright photo of the {}.',\n",
        "    'a dark photo of {article} {}.',\n",
        "    'a dark photo of the {}.',\n",
        "\n",
        "    'a photo of a hard to see {}.',\n",
        "    'a photo of the hard to see {}.',\n",
        "    'a low resolution photo of {article} {}.',\n",
        "    'a low resolution photo of the {}.',\n",
        "    'a cropped photo of {article} {}.',\n",
        "    'a cropped photo of the {}.',\n",
        "    'a close-up photo of {article} {}.',\n",
        "    'a close-up photo of the {}.',\n",
        "    'a jpeg corrupted photo of {article} {}.',\n",
        "    'a jpeg corrupted photo of the {}.',\n",
        "    'a blurry photo of {article} {}.',\n",
        "    'a blurry photo of the {}.',\n",
        "    'a pixelated photo of {article} {}.',\n",
        "    'a pixelated photo of the {}.',\n",
        "\n",
        "    'a black and white photo of the {}.',\n",
        "    'a black and white photo of {article} {}.',\n",
        "\n",
        "    'a plastic {}.',\n",
        "    'the plastic {}.',\n",
        "\n",
        "    'a toy {}.',\n",
        "    'the toy {}.',\n",
        "    'a plushie {}.',\n",
        "    'the plushie {}.',\n",
        "    'a cartoon {}.',\n",
        "    'the cartoon {}.',\n",
        "\n",
        "    'an embroidered {}.',\n",
        "    'the embroidered {}.',\n",
        "\n",
        "    'a painting of the {}.',\n",
        "    'a painting of a {}.',\n",
        "]\n",
        "\n",
        "def build_text_embedding(categories):\n",
        "  if FLAGS.prompt_engineering:\n",
        "    templates = multiple_templates\n",
        "  else:\n",
        "    templates = single_template\n",
        "\n",
        "  run_on_gpu = torch.cuda.is_available()\n",
        "\n",
        "  with torch.no_grad():\n",
        "    all_text_embeddings = []\n",
        "    print(\"Building text embeddings...\")\n",
        "    for category in tqdm(categories):\n",
        "      texts = [\n",
        "        template.format(processed_name(category[\"name\"], rm_dot=True),\n",
        "                        article=article(category[\"name\"]))\n",
        "        for template in templates]\n",
        "      if FLAGS.this_is:\n",
        "        texts = [\n",
        "                 \"This is \" + text if text.startswith(\"a\") or text.startswith(\"the\") else text \n",
        "                 for text in texts\n",
        "                 ]\n",
        "      texts = clip.tokenize(texts) #tokenize\n",
        "      if run_on_gpu:\n",
        "        texts = texts.cuda()\n",
        "      text_embeddings = clip_model.encode_text(texts) #embed with text encoder\n",
        "      text_embeddings /= text_embeddings.norm(dim=-1, keepdim=True)\n",
        "      text_embedding = text_embeddings.mean(dim=0)\n",
        "      text_embedding /= text_embedding.norm()\n",
        "      all_text_embeddings.append(text_embedding)\n",
        "    all_text_embeddings = torch.stack(all_text_embeddings, dim=1)\n",
        "    if run_on_gpu:\n",
        "      all_text_embeddings = all_text_embeddings.cuda()\n",
        "  return all_text_embeddings.cpu().numpy().T"
      ],
      "metadata": {
        "cellView": "form",
        "id": "TOtCs2B9iDMJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Load ViLD model.\n",
        "\n",
        "gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.2)\n",
        "session = tf.Session(graph=tf.Graph(), config=tf.ConfigProto(gpu_options=gpu_options))\n",
        "saved_model_dir = \"./image_path_v2\"\n",
        "_ = tf.saved_model.loader.load(session, [\"serve\"], saved_model_dir)\n",
        "\n",
        "numbered_categories = [{\"name\": str(idx), \"id\": idx,} for idx in range(50)]\n",
        "numbered_category_indices = {cat[\"id\"]: cat for cat in numbered_categories}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "FK5lnSQEiM4r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Non-maximum suppression (NMS).\n",
        "def nms(dets, scores, thresh, max_dets=1000):\n",
        "  \"\"\"Non-maximum suppression.\n",
        "  Args:\n",
        "    dets: [N, 4]\n",
        "    scores: [N,]\n",
        "    thresh: iou threshold. Float\n",
        "    max_dets: int.\n",
        "  \"\"\"\n",
        "  y1 = dets[:, 0]\n",
        "  x1 = dets[:, 1]\n",
        "  y2 = dets[:, 2]\n",
        "  x2 = dets[:, 3]\n",
        "\n",
        "  areas = (x2 - x1) * (y2 - y1)\n",
        "  order = scores.argsort()[::-1]\n",
        "\n",
        "  keep = []\n",
        "  while order.size > 0 and len(keep) < max_dets:\n",
        "    i = order[0]\n",
        "    keep.append(i)\n",
        "\n",
        "    xx1 = np.maximum(x1[i], x1[order[1:]])\n",
        "    yy1 = np.maximum(y1[i], y1[order[1:]])\n",
        "    xx2 = np.minimum(x2[i], x2[order[1:]])\n",
        "    yy2 = np.minimum(y2[i], y2[order[1:]])\n",
        "\n",
        "    w = np.maximum(0.0, xx2 - xx1)\n",
        "    h = np.maximum(0.0, yy2 - yy1)\n",
        "    intersection = w * h\n",
        "    overlap = intersection / (areas[i] + areas[order[1:]] - intersection + 1e-12)\n",
        "\n",
        "    inds = np.where(overlap <= thresh)[0]\n",
        "    order = order[inds + 1]\n",
        "  return keep"
      ],
      "metadata": {
        "cellView": "form",
        "id": "gazf5PBkkzlK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ViLD Result Visualization\n",
        "import PIL.ImageColor as ImageColor\n",
        "import PIL.ImageDraw as ImageDraw\n",
        "import PIL.ImageFont as ImageFont\n",
        "\n",
        "STANDARD_COLORS = [\"White\"]\n",
        "# STANDARD_COLORS = [\n",
        "#     \"AliceBlue\", \"Chartreuse\", \"Aqua\", \"Aquamarine\", \"Azure\", \"Beige\", \"Bisque\",\n",
        "#     \"BlanchedAlmond\", \"BlueViolet\", \"BurlyWood\", \"CadetBlue\", \"AntiqueWhite\",\n",
        "#     \"Chocolate\", \"Coral\", \"CornflowerBlue\", \"Cornsilk\", \"Cyan\",\n",
        "#     \"DarkCyan\", \"DarkGoldenRod\", \"DarkGrey\", \"DarkKhaki\", \"DarkOrange\",\n",
        "#     \"DarkOrchid\", \"DarkSalmon\", \"DarkSeaGreen\", \"DarkTurquoise\", \"DarkViolet\",\n",
        "#     \"DeepPink\", \"DeepSkyBlue\", \"DodgerBlue\", \"FloralWhite\",\n",
        "#     \"ForestGreen\", \"Fuchsia\", \"Gainsboro\", \"GhostWhite\", \"Gold\", \"GoldenRod\",\n",
        "#     \"Salmon\", \"Tan\", \"HoneyDew\", \"HotPink\", \"Ivory\", \"Khaki\",\n",
        "#     \"Lavender\", \"LavenderBlush\", \"LawnGreen\", \"LemonChiffon\", \"LightBlue\",\n",
        "#     \"LightCoral\", \"LightCyan\", \"LightGoldenRodYellow\", \"LightGray\", \"LightGrey\",\n",
        "#     \"LightGreen\", \"LightPink\", \"LightSalmon\", \"LightSeaGreen\", \"LightSkyBlue\",\n",
        "#     \"LightSlateGray\", \"LightSlateGrey\", \"LightSteelBlue\", \"LightYellow\", \"Lime\",\n",
        "#     \"LimeGreen\", \"Linen\", \"Magenta\", \"MediumAquaMarine\", \"MediumOrchid\",\n",
        "#     \"MediumPurple\", \"MediumSeaGreen\", \"MediumSlateBlue\", \"MediumSpringGreen\",\n",
        "#     \"MediumTurquoise\", \"MediumVioletRed\", \"MintCream\", \"MistyRose\", \"Moccasin\",\n",
        "#     \"NavajoWhite\", \"OldLace\", \"Olive\", \"OliveDrab\", \"Orange\",\n",
        "#     \"Orchid\", \"PaleGoldenRod\", \"PaleGreen\", \"PaleTurquoise\", \"PaleVioletRed\",\n",
        "#     \"PapayaWhip\", \"PeachPuff\", \"Peru\", \"Pink\", \"Plum\", \"PowderBlue\", \"Purple\",\n",
        "#     \"RosyBrown\", \"RoyalBlue\", \"SaddleBrown\", \"Green\", \"SandyBrown\",\n",
        "#     \"SeaGreen\", \"SeaShell\", \"Sienna\", \"Silver\", \"SkyBlue\", \"SlateBlue\",\n",
        "#     \"SlateGray\", \"SlateGrey\", \"Snow\", \"SpringGreen\", \"SteelBlue\", \"GreenYellow\",\n",
        "#     \"Teal\", \"Thistle\", \"Tomato\", \"Turquoise\", \"Violet\", \"Wheat\", \"White\",\n",
        "#     \"WhiteSmoke\", \"Yellow\", \"YellowGreen\"\n",
        "# ]\n",
        "\n",
        "def draw_bounding_box_on_image(image,\n",
        "                               ymin,\n",
        "                               xmin,\n",
        "                               ymax,\n",
        "                               xmax,\n",
        "                               color=\"red\",\n",
        "                               thickness=4,\n",
        "                               display_str_list=(),\n",
        "                               use_normalized_coordinates=True):\n",
        "  \"\"\"Adds a bounding box to an image.\n",
        "\n",
        "  Bounding box coordinates can be specified in either absolute (pixel) or\n",
        "  normalized coordinates by setting the use_normalized_coordinates argument.\n",
        "\n",
        "  Each string in display_str_list is displayed on a separate line above the\n",
        "  bounding box in black text on a rectangle filled with the input \"color\".\n",
        "  If the top of the bounding box extends to the edge of the image, the strings\n",
        "  are displayed below the bounding box.\n",
        "\n",
        "  Args:\n",
        "    image: a PIL.Image object.\n",
        "    ymin: ymin of bounding box.\n",
        "    xmin: xmin of bounding box.\n",
        "    ymax: ymax of bounding box.\n",
        "    xmax: xmax of bounding box.\n",
        "    color: color to draw bounding box. Default is red.\n",
        "    thickness: line thickness. Default value is 4.\n",
        "    display_str_list: list of strings to display in box\n",
        "                      (each to be shown on its own line).\n",
        "    use_normalized_coordinates: If True (default), treat coordinates\n",
        "      ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
        "      coordinates as absolute.\n",
        "  \"\"\"\n",
        "  draw = ImageDraw.Draw(image)\n",
        "  im_width, im_height = image.size\n",
        "  if use_normalized_coordinates:\n",
        "    (left, right, top, bottom) = (xmin * im_width, xmax * im_width,\n",
        "                                  ymin * im_height, ymax * im_height)\n",
        "  else:\n",
        "    (left, right, top, bottom) = (xmin, xmax, ymin, ymax)\n",
        "  draw.line([(left, top), (left, bottom), (right, bottom),\n",
        "             (right, top), (left, top)], width=thickness, fill=color)\n",
        "  try:\n",
        "    font = ImageFont.truetype(\"arial.ttf\", 24)\n",
        "  except IOError:\n",
        "    font = ImageFont.load_default()\n",
        "\n",
        "  # If the total height of the display strings added to the top of the bounding\n",
        "  # box exceeds the top of the image, stack the strings below the bounding box\n",
        "  # instead of above.\n",
        "  display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]\n",
        "  # Each display_str has a top and bottom margin of 0.05x.\n",
        "  total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)\n",
        "\n",
        "  if top > total_display_str_height:\n",
        "    text_bottom = top\n",
        "  else:\n",
        "    text_bottom = bottom + total_display_str_height\n",
        "  # Reverse list and print from bottom to top.\n",
        "  for display_str in display_str_list[::-1]:\n",
        "    text_left = min(5, left)\n",
        "    text_width, text_height = font.getsize(display_str)\n",
        "    margin = np.ceil(0.05 * text_height)\n",
        "    draw.rectangle(\n",
        "        [(left, text_bottom - text_height - 2 * margin), (left + text_width,\n",
        "                                                          text_bottom)],\n",
        "        fill=color)\n",
        "    draw.text(\n",
        "        (left + margin, text_bottom - text_height - margin),\n",
        "        display_str,\n",
        "        fill=\"black\",\n",
        "        font=font)\n",
        "    text_bottom -= text_height - 2 * margin\n",
        "\n",
        "def draw_bounding_box_on_image_array(image,\n",
        "                                     ymin,\n",
        "                                     xmin,\n",
        "                                     ymax,\n",
        "                                     xmax,\n",
        "                                     color=\"red\",\n",
        "                                     thickness=4,\n",
        "                                     display_str_list=(),\n",
        "                                     use_normalized_coordinates=True):\n",
        "  \"\"\"Adds a bounding box to an image (numpy array).\n",
        "\n",
        "  Bounding box coordinates can be specified in either absolute (pixel) or\n",
        "  normalized coordinates by setting the use_normalized_coordinates argument.\n",
        "\n",
        "  Args:\n",
        "    image: a numpy array with shape [height, width, 3].\n",
        "    ymin: ymin of bounding box.\n",
        "    xmin: xmin of bounding box.\n",
        "    ymax: ymax of bounding box.\n",
        "    xmax: xmax of bounding box.\n",
        "    color: color to draw bounding box. Default is red.\n",
        "    thickness: line thickness. Default value is 4.\n",
        "    display_str_list: list of strings to display in box\n",
        "                      (each to be shown on its own line).\n",
        "    use_normalized_coordinates: If True (default), treat coordinates\n",
        "      ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
        "      coordinates as absolute.\n",
        "  \"\"\"\n",
        "  image_pil = Image.fromarray(np.uint8(image)).convert(\"RGB\")\n",
        "  draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,\n",
        "                             thickness, display_str_list,\n",
        "                             use_normalized_coordinates)\n",
        "  np.copyto(image, np.array(image_pil))\n",
        "\n",
        "\n",
        "def draw_mask_on_image_array(image, mask, color=\"red\", alpha=0.4):\n",
        "  \"\"\"Draws mask on an image.\n",
        "\n",
        "  Args:\n",
        "    image: uint8 numpy array with shape (img_height, img_height, 3)\n",
        "    mask: a uint8 numpy array of shape (img_height, img_height) with\n",
        "      values between either 0 or 1.\n",
        "    color: color to draw the keypoints with. Default is red.\n",
        "    alpha: transparency value between 0 and 1. (default: 0.4)\n",
        "\n",
        "  Raises:\n",
        "    ValueError: On incorrect data type for image or masks.\n",
        "  \"\"\"\n",
        "  if image.dtype != np.uint8:\n",
        "    raise ValueError(\"`image` not of type np.uint8\")\n",
        "  if mask.dtype != np.uint8:\n",
        "    raise ValueError(\"`mask` not of type np.uint8\")\n",
        "  if np.any(np.logical_and(mask != 1, mask != 0)):\n",
        "    raise ValueError(\"`mask` elements should be in [0, 1]\")\n",
        "  if image.shape[:2] != mask.shape:\n",
        "    raise ValueError(\"The image has spatial dimensions %s but the mask has \"\n",
        "                     \"dimensions %s\" % (image.shape[:2], mask.shape))\n",
        "  rgb = ImageColor.getrgb(color)\n",
        "  pil_image = Image.fromarray(image)\n",
        "\n",
        "  solid_color = np.expand_dims(\n",
        "      np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])\n",
        "  pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert(\"RGBA\")\n",
        "  pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert(\"L\")\n",
        "  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)\n",
        "  np.copyto(image, np.array(pil_image.convert(\"RGB\")))\n",
        "\n",
        "def visualize_boxes_and_labels_on_image_array(\n",
        "    image,\n",
        "    boxes,\n",
        "    classes,\n",
        "    scores,\n",
        "    category_index,\n",
        "    instance_masks=None,\n",
        "    instance_boundaries=None,\n",
        "    use_normalized_coordinates=False,\n",
        "    max_boxes_to_draw=20,\n",
        "    min_score_thresh=.5,\n",
        "    agnostic_mode=False,\n",
        "    line_thickness=1,\n",
        "    groundtruth_box_visualization_color=\"black\",\n",
        "    skip_scores=False,\n",
        "    skip_labels=False,\n",
        "    mask_alpha=0.4,\n",
        "    plot_color=None,\n",
        "):\n",
        "  \"\"\"Overlay labeled boxes on an image with formatted scores and label names.\n",
        "\n",
        "  This function groups boxes that correspond to the same location\n",
        "  and creates a display string for each detection and overlays these\n",
        "  on the image. Note that this function modifies the image in place, and returns\n",
        "  that same image.\n",
        "\n",
        "  Args:\n",
        "    image: uint8 numpy array with shape (img_height, img_width, 3)\n",
        "    boxes: a numpy array of shape [N, 4]\n",
        "    classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
        "      and match the keys in the label map.\n",
        "    scores: a numpy array of shape [N] or None.  If scores=None, then\n",
        "      this function assumes that the boxes to be plotted are groundtruth\n",
        "      boxes and plot all boxes as black with no classes or scores.\n",
        "    category_index: a dict containing category dictionaries (each holding\n",
        "      category index `id` and category name `name`) keyed by category indices.\n",
        "    instance_masks: a numpy array of shape [N, image_height, image_width] with\n",
        "      values ranging between 0 and 1, can be None.\n",
        "    instance_boundaries: a numpy array of shape [N, image_height, image_width]\n",
        "      with values ranging between 0 and 1, can be None.\n",
        "    use_normalized_coordinates: whether boxes is to be interpreted as\n",
        "      normalized coordinates or not.\n",
        "    max_boxes_to_draw: maximum number of boxes to visualize.  If None, draw\n",
        "      all boxes.\n",
        "    min_score_thresh: minimum score threshold for a box to be visualized\n",
        "    agnostic_mode: boolean (default: False) controlling whether to evaluate in\n",
        "      class-agnostic mode or not.  This mode will display scores but ignore\n",
        "      classes.\n",
        "    line_thickness: integer (default: 4) controlling line width of the boxes.\n",
        "    groundtruth_box_visualization_color: box color for visualizing groundtruth\n",
        "      boxes\n",
        "    skip_scores: whether to skip score when drawing a single detection\n",
        "    skip_labels: whether to skip label when drawing a single detection\n",
        "\n",
        "  Returns:\n",
        "    uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.\n",
        "  \"\"\"\n",
        "  # Create a display string (and color) for every box location, group any boxes\n",
        "  # that correspond to the same location.\n",
        "  box_to_display_str_map = collections.defaultdict(list)\n",
        "  box_to_color_map = collections.defaultdict(str)\n",
        "  box_to_instance_masks_map = {}\n",
        "  box_to_score_map = {}\n",
        "  box_to_instance_boundaries_map = {}\n",
        "  \n",
        "  if not max_boxes_to_draw:\n",
        "    max_boxes_to_draw = boxes.shape[0]\n",
        "  for i in range(min(max_boxes_to_draw, boxes.shape[0])):\n",
        "    if scores is None or scores[i] > min_score_thresh:\n",
        "      box = tuple(boxes[i].tolist())\n",
        "      if instance_masks is not None:\n",
        "        box_to_instance_masks_map[box] = instance_masks[i]\n",
        "      if instance_boundaries is not None:\n",
        "        box_to_instance_boundaries_map[box] = instance_boundaries[i]\n",
        "      if scores is None:\n",
        "        box_to_color_map[box] = groundtruth_box_visualization_color\n",
        "      else:\n",
        "        display_str = \"\"\n",
        "        if not skip_labels:\n",
        "          if not agnostic_mode:\n",
        "            if classes[i] in list(category_index.keys()):\n",
        "              class_name = category_index[classes[i]][\"name\"]\n",
        "            else:\n",
        "              class_name = \"N/A\"\n",
        "            display_str = str(class_name)\n",
        "        if not skip_scores:\n",
        "          if not display_str:\n",
        "            display_str = \"{}%\".format(int(100*scores[i]))\n",
        "          else:\n",
        "            float_score = (\"%.2f\" % scores[i]).lstrip(\"0\")\n",
        "            display_str = \"{}: {}\".format(display_str, float_score)\n",
        "          box_to_score_map[box] = int(100*scores[i])\n",
        "\n",
        "        box_to_display_str_map[box].append(display_str)\n",
        "        if plot_color is not None:\n",
        "          box_to_color_map[box] = plot_color\n",
        "        elif agnostic_mode:\n",
        "          box_to_color_map[box] = \"DarkOrange\"\n",
        "        else:\n",
        "          box_to_color_map[box] = STANDARD_COLORS[\n",
        "              classes[i] % len(STANDARD_COLORS)]\n",
        "\n",
        "  # Handle the case when box_to_score_map is empty.\n",
        "  if box_to_score_map:\n",
        "    box_color_iter = sorted(\n",
        "        box_to_color_map.items(), key=lambda kv: box_to_score_map[kv[0]])\n",
        "  else:\n",
        "    box_color_iter = box_to_color_map.items()\n",
        "\n",
        "  # Draw all boxes onto image.\n",
        "  for box, color in box_color_iter:\n",
        "    ymin, xmin, ymax, xmax = box\n",
        "    if instance_masks is not None:\n",
        "      draw_mask_on_image_array(\n",
        "          image,\n",
        "          box_to_instance_masks_map[box],\n",
        "          color=color,\n",
        "          alpha=mask_alpha\n",
        "      )\n",
        "    if instance_boundaries is not None:\n",
        "      draw_mask_on_image_array(\n",
        "          image,\n",
        "          box_to_instance_boundaries_map[box],\n",
        "          color=\"red\",\n",
        "          alpha=1.0\n",
        "      )\n",
        "    draw_bounding_box_on_image_array(\n",
        "        image,\n",
        "        ymin,\n",
        "        xmin,\n",
        "        ymax,\n",
        "        xmax,\n",
        "        color=color,\n",
        "        thickness=line_thickness,\n",
        "        display_str_list=box_to_display_str_map[box],\n",
        "        use_normalized_coordinates=use_normalized_coordinates)\n",
        "    \n",
        "  return image\n",
        "\n",
        "\n",
        "def paste_instance_masks(masks,\n",
        "                         detected_boxes,\n",
        "                         image_height,\n",
        "                         image_width):\n",
        "  \"\"\"Paste instance masks to generate the image segmentation results.\n",
        "\n",
        "  Args:\n",
        "    masks: a numpy array of shape [N, mask_height, mask_width] representing the\n",
        "      instance masks w.r.t. the `detected_boxes`.\n",
        "    detected_boxes: a numpy array of shape [N, 4] representing the reference\n",
        "      bounding boxes.\n",
        "    image_height: an integer representing the height of the image.\n",
        "    image_width: an integer representing the width of the image.\n",
        "\n",
        "  Returns:\n",
        "    segms: a numpy array of shape [N, image_height, image_width] representing\n",
        "      the instance masks *pasted* on the image canvas.\n",
        "  \"\"\"\n",
        "\n",
        "  def expand_boxes(boxes, scale):\n",
        "    \"\"\"Expands an array of boxes by a given scale.\"\"\"\n",
        "    # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/boxes.py#L227  # pylint: disable=line-too-long\n",
        "    # The `boxes` in the reference implementation is in [x1, y1, x2, y2] form,\n",
        "    # whereas `boxes` here is in [x1, y1, w, h] form\n",
        "    w_half = boxes[:, 2] * .5\n",
        "    h_half = boxes[:, 3] * .5\n",
        "    x_c = boxes[:, 0] + w_half\n",
        "    y_c = boxes[:, 1] + h_half\n",
        "\n",
        "    w_half *= scale\n",
        "    h_half *= scale\n",
        "\n",
        "    boxes_exp = np.zeros(boxes.shape)\n",
        "    boxes_exp[:, 0] = x_c - w_half\n",
        "    boxes_exp[:, 2] = x_c + w_half\n",
        "    boxes_exp[:, 1] = y_c - h_half\n",
        "    boxes_exp[:, 3] = y_c + h_half\n",
        "\n",
        "    return boxes_exp\n",
        "\n",
        "  # Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/core/test.py#L812  # pylint: disable=line-too-long\n",
        "  # To work around an issue with cv2.resize (it seems to automatically pad\n",
        "  # with repeated border values), we manually zero-pad the masks by 1 pixel\n",
        "  # prior to resizing back to the original image resolution. This prevents\n",
        "  # \"top hat\" artifacts. We therefore need to expand the reference boxes by an\n",
        "  # appropriate factor.\n",
        "  _, mask_height, mask_width = masks.shape\n",
        "  scale = max((mask_width + 2.0) / mask_width,\n",
        "              (mask_height + 2.0) / mask_height)\n",
        "\n",
        "  ref_boxes = expand_boxes(detected_boxes, scale)\n",
        "  ref_boxes = ref_boxes.astype(np.int32)\n",
        "  padded_mask = np.zeros((mask_height + 2, mask_width + 2), dtype=np.float32)\n",
        "  segms = []\n",
        "  for mask_ind, mask in enumerate(masks):\n",
        "    im_mask = np.zeros((image_height, image_width), dtype=np.uint8)\n",
        "    # Process mask inside bounding boxes.\n",
        "    padded_mask[1:-1, 1:-1] = mask[:, :]\n",
        "\n",
        "    ref_box = ref_boxes[mask_ind, :]\n",
        "    w = ref_box[2] - ref_box[0] + 1\n",
        "    h = ref_box[3] - ref_box[1] + 1\n",
        "    w = np.maximum(w, 1)\n",
        "    h = np.maximum(h, 1)\n",
        "\n",
        "    mask = cv2.resize(padded_mask, (w, h))\n",
        "    mask = np.array(mask > 0.5, dtype=np.uint8)\n",
        "\n",
        "    x_0 = min(max(ref_box[0], 0), image_width)\n",
        "    x_1 = min(max(ref_box[2] + 1, 0), image_width)\n",
        "    y_0 = min(max(ref_box[1], 0), image_height)\n",
        "    y_1 = min(max(ref_box[3] + 1, 0), image_height)\n",
        "\n",
        "    im_mask[y_0:y_1, x_0:x_1] = mask[\n",
        "        (y_0 - ref_box[1]):(y_1 - ref_box[1]),\n",
        "        (x_0 - ref_box[0]):(x_1 - ref_box[0])\n",
        "    ]\n",
        "    segms.append(im_mask)\n",
        "\n",
        "  segms = np.array(segms)\n",
        "  assert masks.shape[0] == segms.shape[0]\n",
        "  return segms  "
      ],
      "metadata": {
        "cellView": "form",
        "id": "JO63v4Xuk1N6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Plot instance masks.\n",
        "def plot_mask(color, alpha, original_image, mask):\n",
        "  rgb = ImageColor.getrgb(color)\n",
        "  pil_image = Image.fromarray(original_image)\n",
        "\n",
        "  solid_color = np.expand_dims(\n",
        "      np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])\n",
        "  pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert(\"RGBA\")\n",
        "  pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert(\"L\")\n",
        "  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)\n",
        "  img_w_mask = np.array(pil_image.convert(\"RGB\"))\n",
        "  return img_w_mask\n",
        "\n",
        "%matplotlib inline\n",
        "def display_image(path_or_array, size=(10, 10)):\n",
        "  if isinstance(path_or_array, str):\n",
        "    image = np.asarray(Image.open(open(image_path, \"rb\")).convert(\"RGB\"))\n",
        "  else:\n",
        "    image = path_or_array\n",
        "  \n",
        "  plt.figure(figsize=size)\n",
        "  plt.imshow(image)\n",
        "  plt.axis(\"off\")\n",
        "  plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "kKKH_wTak3F_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Define ViLD forward pass.\n",
        "\n",
        "def vild(image_path, category_name_string, params, plot_on=True, prompt_swaps=[]):\n",
        "  #################################################################\n",
        "  # Preprocessing categories and get params\n",
        "  for a, b in prompt_swaps:\n",
        "    category_name_string = category_name_string.replace(a, b)\n",
        "  category_names = [x.strip() for x in category_name_string.split(\";\")]\n",
        "  category_names = [\"background\"] + category_names\n",
        "  categories = [{\"name\": item, \"id\": idx+1,} for idx, item in enumerate(category_names)]\n",
        "  category_indices = {cat[\"id\"]: cat for cat in categories}\n",
        "  \n",
        "  max_boxes_to_draw, nms_threshold, min_rpn_score_thresh, min_box_area, max_box_area = params\n",
        "  fig_size_h = min(max(5, int(len(category_names) / 2.5) ), 10)\n",
        "\n",
        "\n",
        "  #################################################################\n",
        "  # Obtain results and read image\n",
        "  roi_boxes, roi_scores, detection_boxes, scores_unused, box_outputs, detection_masks, visual_features, image_info = session.run(\n",
        "        [\"RoiBoxes:0\", \"RoiScores:0\", \"2ndStageBoxes:0\", \"2ndStageScoresUnused:0\", \"BoxOutputs:0\", \"MaskOutputs:0\", \"VisualFeatOutputs:0\", \"ImageInfo:0\"],\n",
        "        feed_dict={\"Placeholder:0\": [image_path,]})\n",
        "  \n",
        "  roi_boxes = np.squeeze(roi_boxes, axis=0)  # squeeze\n",
        "  # no need to clip the boxes, already done\n",
        "  roi_scores = np.squeeze(roi_scores, axis=0)\n",
        "\n",
        "  detection_boxes = np.squeeze(detection_boxes, axis=(0, 2))\n",
        "  scores_unused = np.squeeze(scores_unused, axis=0)\n",
        "  box_outputs = np.squeeze(box_outputs, axis=0)\n",
        "  detection_masks = np.squeeze(detection_masks, axis=0)\n",
        "  visual_features = np.squeeze(visual_features, axis=0)\n",
        "\n",
        "  image_info = np.squeeze(image_info, axis=0)  # obtain image info\n",
        "  image_scale = np.tile(image_info[2:3, :], (1, 2))\n",
        "  image_height = int(image_info[0, 0])\n",
        "  image_width = int(image_info[0, 1])\n",
        "\n",
        "  rescaled_detection_boxes = detection_boxes / image_scale # rescale\n",
        "\n",
        "  # Read image\n",
        "  image = np.asarray(Image.open(open(image_path, \"rb\")).convert(\"RGB\"))\n",
        "  assert image_height == image.shape[0]\n",
        "  assert image_width == image.shape[1]\n",
        "\n",
        "\n",
        "  #################################################################\n",
        "  # Filter boxes\n",
        "\n",
        "  # Apply non-maximum suppression to detected boxes with nms threshold.\n",
        "  nmsed_indices = nms(\n",
        "      detection_boxes,\n",
        "      roi_scores,\n",
        "      thresh=nms_threshold\n",
        "      )\n",
        "\n",
        "  # Compute RPN box size.\n",
        "  box_sizes = (rescaled_detection_boxes[:, 2] - rescaled_detection_boxes[:, 0]) * (rescaled_detection_boxes[:, 3] - rescaled_detection_boxes[:, 1])\n",
        "\n",
        "  # Filter out invalid rois (nmsed rois)\n",
        "  valid_indices = np.where(\n",
        "      np.logical_and(\n",
        "        np.isin(np.arange(len(roi_scores), dtype=np.int), nmsed_indices),\n",
        "        np.logical_and(\n",
        "            np.logical_not(np.all(roi_boxes == 0., axis=-1)),\n",
        "            np.logical_and(\n",
        "              roi_scores >= min_rpn_score_thresh,\n",
        "              np.logical_and(\n",
        "                box_sizes > min_box_area,\n",
        "                box_sizes < max_box_area\n",
        "                )\n",
        "              )\n",
        "        )    \n",
        "      )\n",
        "  )[0]\n",
        "\n",
        "  detection_roi_scores = roi_scores[valid_indices][:max_boxes_to_draw, ...]\n",
        "  detection_boxes = detection_boxes[valid_indices][:max_boxes_to_draw, ...]\n",
        "  detection_masks = detection_masks[valid_indices][:max_boxes_to_draw, ...]\n",
        "  detection_visual_feat = visual_features[valid_indices][:max_boxes_to_draw, ...]\n",
        "  rescaled_detection_boxes = rescaled_detection_boxes[valid_indices][:max_boxes_to_draw, ...]\n",
        "\n",
        "\n",
        "  #################################################################\n",
        "  # Compute text embeddings and detection scores, and rank results\n",
        "  text_features = build_text_embedding(categories)\n",
        "  \n",
        "  raw_scores = detection_visual_feat.dot(text_features.T)\n",
        "  if FLAGS.use_softmax:\n",
        "    scores_all = softmax(FLAGS.temperature * raw_scores, axis=-1)\n",
        "  else:\n",
        "    scores_all = raw_scores\n",
        "\n",
        "  indices = np.argsort(-np.max(scores_all, axis=1))  # Results are ranked by scores\n",
        "  indices_fg = np.array([i for i in indices if np.argmax(scores_all[i]) != 0])\n",
        "\n",
        "  \n",
        "  #################################################################\n",
        "  # Print found_objects\n",
        "  found_objects = []\n",
        "  for a, b in prompt_swaps:\n",
        "    category_names = [name.replace(b, a) for name in category_names]  # Extra prompt engineering.\n",
        "  for anno_idx in indices[0:int(rescaled_detection_boxes.shape[0])]:\n",
        "    scores = scores_all[anno_idx]\n",
        "    if np.argmax(scores) == 0:\n",
        "      continue\n",
        "    found_object = category_names[np.argmax(scores)]\n",
        "    if found_object == \"background\":\n",
        "      continue\n",
        "    print(\"Found a\", found_object, \"with score:\", np.max(scores))\n",
        "    found_objects.append(category_names[np.argmax(scores)])\n",
        "  if not plot_on:\n",
        "    return found_objects\n",
        "  \n",
        "\n",
        "  #################################################################\n",
        "  # Plot detected boxes on the input image.\n",
        "  ymin, xmin, ymax, xmax = np.split(rescaled_detection_boxes, 4, axis=-1)\n",
        "  processed_boxes = np.concatenate([xmin, ymin, xmax - xmin, ymax - ymin], axis=-1)\n",
        "  segmentations = paste_instance_masks(detection_masks, processed_boxes, image_height, image_width)\n",
        "\n",
        "  if len(indices_fg) == 0:\n",
        "    display_image(np.array(image), size=overall_fig_size)\n",
        "    print(\"ViLD does not detect anything belong to the given category\")\n",
        "\n",
        "  else:\n",
        "    image_with_detections = visualize_boxes_and_labels_on_image_array(\n",
        "        np.array(image),\n",
        "        rescaled_detection_boxes[indices_fg],\n",
        "        valid_indices[:max_boxes_to_draw][indices_fg],\n",
        "        detection_roi_scores[indices_fg],    \n",
        "        numbered_category_indices,\n",
        "        instance_masks=segmentations[indices_fg],\n",
        "        use_normalized_coordinates=False,\n",
        "        max_boxes_to_draw=max_boxes_to_draw,\n",
        "        min_score_thresh=min_rpn_score_thresh,\n",
        "        skip_scores=False,\n",
        "        skip_labels=True)\n",
        "\n",
        "    # plt.figure(figsize=overall_fig_size)\n",
        "    plt.imshow(image_with_detections)\n",
        "    # plt.axis(\"off\")\n",
        "    plt.title(\"ViLD detected objects and RPN scores.\")\n",
        "    plt.show()\n",
        "\n",
        "  return found_objects"
      ],
      "metadata": {
        "cellView": "form",
        "id": "X2-SsyEolAr8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "category_names = ['blue block',\n",
        "                  'red block',\n",
        "                  'green block',\n",
        "                  'orange block',\n",
        "                  'yellow block',\n",
        "                  'purple block',\n",
        "                  'pink block',\n",
        "                  'cyan block',\n",
        "                  'brown block',\n",
        "                  'gray block',\n",
        "\n",
        "                  'blue bowl',\n",
        "                  'red bowl',\n",
        "                  'green bowl',\n",
        "                  'orange bowl',\n",
        "                  'yellow bowl',\n",
        "                  'purple bowl',\n",
        "                  'pink bowl',\n",
        "                  'cyan bowl',\n",
        "                  'brown bowl',\n",
        "                  'gray bowl']\n",
        "image_path = 'tmp.jpg'\n",
        "\n",
        "#@markdown ViLD settings.\n",
        "category_name_string = \";\".join(category_names)\n",
        "max_boxes_to_draw = 8 #@param {type:\"integer\"}\n",
        "\n",
        "# Extra prompt engineering: swap A with B for every (A, B) in list.\n",
        "prompt_swaps = [('block', 'cube')]\n",
        "\n",
        "nms_threshold = 0.4 #@param {type:\"slider\", min:0, max:0.9, step:0.05}\n",
        "min_rpn_score_thresh = 0.4  #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "min_box_area = 10 #@param {type:\"slider\", min:0, max:10000, step:1.0}\n",
        "max_box_area = 3000  #@param {type:\"slider\", min:0, max:10000, step:1.0}\n",
        "vild_params = max_boxes_to_draw, nms_threshold, min_rpn_score_thresh, min_box_area, max_box_area\n",
        "found_objects = vild(image_path, category_name_string, vild_params, plot_on=True, prompt_swaps=prompt_swaps)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8jwAERXZl2Yl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Scripted Expert**\n",
        "Scripted pick & place oracle to collect expert demonstrations."
      ],
      "metadata": {
        "id": "K9K0Bxh2Gj8p"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ScriptedPolicy():\n",
        "\n",
        "  def __init__(self, env):\n",
        "    self.env = env\n",
        "\n",
        "  def step(self, text, obs):\n",
        "    print(f'Input: {text}')\n",
        "\n",
        "    # Parse pick and place targets.\n",
        "    pick_text, place_text = text.split('and')\n",
        "    pick_target, place_target = None, None\n",
        "    for name in PICK_TARGETS.keys():\n",
        "      if name in pick_text:\n",
        "        pick_target = name\n",
        "        break\n",
        "    for name in PLACE_TARGETS.keys():\n",
        "      if name in place_text:\n",
        "        place_target = name\n",
        "        break\n",
        "\n",
        "    # Admissable targets only.\n",
        "    assert pick_target is not None\n",
        "    assert place_target is not None\n",
        "\n",
        "    pick_id = self.env.obj_name_to_id[pick_target]\n",
        "    pick_pose = pybullet.getBasePositionAndOrientation(pick_id)\n",
        "    pick_position = np.float32(pick_pose[0])\n",
        "\n",
        "    if place_target in self.env.obj_name_to_id:\n",
        "      place_id = self.env.obj_name_to_id[place_target]\n",
        "      place_pose = pybullet.getBasePositionAndOrientation(place_id)\n",
        "      place_position = np.float32(place_pose[0])\n",
        "    else:\n",
        "      place_position = np.float32(PLACE_TARGETS[place_target])\n",
        "\n",
        "    # Add some noise to pick and place positions.\n",
        "    # pick_position[:2] += np.random.normal(scale=0.01)\n",
        "    place_position[:2] += np.random.normal(scale=0.01)\n",
        "\n",
        "    act = {'pick': pick_position, 'place': place_position}\n",
        "    return act"
      ],
      "metadata": {
        "id": "QYwloVzrHBbE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Dataset**"
      ],
      "metadata": {
        "id": "B7JYho2gHNES"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Collect demonstrations with a scripted expert, or download a pre-generated dataset.\n",
        "load_pregenerated = True  #@param {type:\"boolean\"}\n",
        "\n",
        "# Load pre-existing dataset.\n",
        "if load_pregenerated:\n",
        "  if not os.path.exists('dataset-9999.pkl'):\n",
        "    # !gdown --id 1TECwTIfawxkRYbzlAey0z1mqXKcyfPc-\n",
        "    !gdown --id 1yCz6C-6eLWb4SFYKdkM-wz5tlMjbG2h8\n",
        "  dataset = pickle.load(open('dataset-9999.pkl', 'rb'))  # ~10K samples.\n",
        "  dataset_size = len(dataset['text'])\n",
        "\n",
        "# Generate new dataset.\n",
        "else:\n",
        "  dataset = {}\n",
        "  dataset_size = 2  # Size of new dataset.\n",
        "  dataset['image'] = np.zeros((dataset_size, 224, 224, 3), dtype=np.uint8)\n",
        "  dataset['pick_yx'] = np.zeros((dataset_size, 2), dtype=np.int32)\n",
        "  dataset['place_yx'] = np.zeros((dataset_size, 2), dtype=np.int32)\n",
        "  dataset['text'] = []\n",
        "  policy = ScriptedPolicy(env)\n",
        "  data_idx = 0\n",
        "  while data_idx < dataset_size:\n",
        "    np.random.seed(data_idx)\n",
        "    num_pick, num_place = 3, 3\n",
        "\n",
        "    # Select random objects for data collection.\n",
        "    pick_items = list(PICK_TARGETS.keys())\n",
        "    pick_items = np.random.choice(pick_items, size=num_pick, replace=False)\n",
        "    place_items = list(PLACE_TARGETS.keys())\n",
        "    for pick_item in pick_items:  # For simplicity: place items != pick items.\n",
        "      place_items.remove(pick_item)\n",
        "    place_items = np.random.choice(place_items, size=num_place, replace=False)\n",
        "    config = {'pick': pick_items, 'place': place_items}\n",
        "\n",
        "    # Initialize environment with selected objects.\n",
        "    obs = env.reset(config)\n",
        "\n",
        "    # Create text prompts.\n",
        "    prompts = []\n",
        "    for i in range(len(pick_items)):\n",
        "      pick_item = pick_items[i]\n",
        "      place_item = place_items[i]\n",
        "      prompts.append(f'Pick the {pick_item} and place it on the {place_item}.')\n",
        "\n",
        "    # Execute 3 pick and place actions.\n",
        "    for prompt in prompts:\n",
        "      act = policy.step(prompt, obs)\n",
        "      dataset['text'].append(prompt)\n",
        "      dataset['image'][data_idx, ...] = obs['image'].copy()\n",
        "      dataset['pick_yx'][data_idx, ...] = xyz_to_pix(act['pick'])\n",
        "      dataset['place_yx'][data_idx, ...] = xyz_to_pix(act['place'])\n",
        "      data_idx += 1\n",
        "      obs, _, _, _ = env.step(act)\n",
        "      debug_clip = ImageSequenceClip(env.cache_video, fps=25)\n",
        "      display(debug_clip.ipython_display(autoplay=1, loop=1))\n",
        "      env.cache_video = []\n",
        "      if data_idx >= dataset_size:\n",
        "        break\n",
        "\n",
        "  pickle.dump(dataset, open(f'dataset-{dataset_size}.pkl', 'wb'))"
      ],
      "metadata": {
        "id": "4_WG2G1LG2jz",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Show a demonstration example from the dataset.\n",
        "\n",
        "img = dataset['image'][0]\n",
        "pick_yx = dataset['pick_yx'][0]\n",
        "place_yx = dataset['place_yx'][0]\n",
        "text = dataset['text'][0]\n",
        "plt.title(text)\n",
        "plt.imshow(img)\n",
        "plt.arrow(pick_yx[1], pick_yx[0], place_yx[1]-pick_yx[1], place_yx[0]-pick_yx[0], color='w', head_starts_at_zero=False, head_width=7, length_includes_head=True)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "-wVcfLWlYm73",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Model**\n",
        "Define a variant of [CLIPort](https://cliport.github.io/) architecture ([Transporter Nets](https://transporternets.github.io/) + [CLIP](https://arxiv.org/pdf/2103.00020.pdf))"
      ],
      "metadata": {
        "id": "IIxTA71w3ZA0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ResNetBlock(nn.Module):\n",
        "  \"\"\"ResNet pre-Activation block. https://arxiv.org/pdf/1603.05027.pdf\"\"\"\n",
        "  features: int\n",
        "  stride: int = 1\n",
        "\n",
        "  def setup(self):\n",
        "    self.conv0 = nn.Conv(self.features // 4, (1, 1), (self.stride, self.stride))\n",
        "    self.conv1 = nn.Conv(self.features // 4, (3, 3))\n",
        "    self.conv2 = nn.Conv(self.features, (1, 1))\n",
        "    self.conv3 = nn.Conv(self.features, (1, 1), (self.stride, self.stride))\n",
        "\n",
        "  def __call__(self, x):\n",
        "    y = self.conv0(nn.relu(x))\n",
        "    y = self.conv1(nn.relu(y))\n",
        "    y = self.conv2(nn.relu(y))\n",
        "    if x.shape != y.shape:\n",
        "      x = self.conv3(nn.relu(x))\n",
        "    return x + y\n",
        "\n",
        "\n",
        "class UpSample(nn.Module):\n",
        "  \"\"\"Simple 2D 2x bilinear upsample.\"\"\"\n",
        "\n",
        "  def __call__(self, x):\n",
        "    B, H, W, C = x.shape\n",
        "    new_shape = (B, H * 2, W * 2, C)\n",
        "    return jax.image.resize(x, new_shape, 'bilinear')\n",
        "\n",
        "\n",
        "class ResNet(nn.Module):\n",
        "  \"\"\"Hourglass 53-layer ResNet with 8-stride.\"\"\"\n",
        "  out_dim: int\n",
        "\n",
        "  def setup(self):\n",
        "    self.dense0 = nn.Dense(8)\n",
        "\n",
        "    self.conv0 = nn.Conv(64, (3, 3), (1, 1))\n",
        "    self.block0 = ResNetBlock(64)\n",
        "    self.block1 = ResNetBlock(64)\n",
        "    self.block2 = ResNetBlock(128, stride=2)\n",
        "    self.block3 = ResNetBlock(128)\n",
        "    self.block4 = ResNetBlock(256, stride=2)\n",
        "    self.block5 = ResNetBlock(256)\n",
        "    self.block6 = ResNetBlock(512, stride=2)\n",
        "    self.block7 = ResNetBlock(512)\n",
        "\n",
        "    self.block8 = ResNetBlock(256)\n",
        "    self.block9 = ResNetBlock(256)\n",
        "    self.upsample0 = UpSample()\n",
        "    self.block10 = ResNetBlock(128)\n",
        "    self.block11 = ResNetBlock(128)\n",
        "    self.upsample1 = UpSample()\n",
        "    self.block12 = ResNetBlock(64)\n",
        "    self.block13 = ResNetBlock(64)\n",
        "    self.upsample2 = UpSample()\n",
        "    self.block14 = ResNetBlock(16)\n",
        "    self.block15 = ResNetBlock(16)\n",
        "    self.conv1 = nn.Conv(self.out_dim, (3, 3), (1, 1))\n",
        "\n",
        "  def __call__(self, x, text):\n",
        "\n",
        "    # # Project and concatenate CLIP features (early fusion).\n",
        "    # text = self.dense0(text)\n",
        "    # text = jnp.expand_dims(text, axis=(1, 2))\n",
        "    # text = jnp.broadcast_to(text, x.shape[:3] + (8,))\n",
        "    # x = jnp.concatenate((x, text), axis=-1)\n",
        "\n",
        "    x = self.conv0(x)\n",
        "    x = self.block0(x)\n",
        "    x = self.block1(x)\n",
        "    x = self.block2(x)\n",
        "    x = self.block3(x)\n",
        "    x = self.block4(x)\n",
        "    x = self.block5(x)\n",
        "    x = self.block6(x)\n",
        "    x = self.block7(x)\n",
        "\n",
        "    # Concatenate CLIP features (mid-fusion).\n",
        "    text = jnp.expand_dims(text, axis=(1, 2))\n",
        "    text = jnp.broadcast_to(text, x.shape)\n",
        "    x = jnp.concatenate((x, text), axis=-1)\n",
        "\n",
        "    x = self.block8(x)\n",
        "    x = self.block9(x)\n",
        "    x = self.upsample0(x)\n",
        "    x = self.block10(x)\n",
        "    x = self.block11(x)\n",
        "    x = self.upsample1(x)\n",
        "    x = self.block12(x)\n",
        "    x = self.block13(x)\n",
        "    x = self.upsample2(x)\n",
        "    x = self.block14(x)\n",
        "    x = self.block15(x)\n",
        "    x = self.conv1(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "class TransporterNets(nn.Module):\n",
        "  \"\"\"TransporterNet with 3 ResNets (translation only).\"\"\"\n",
        "\n",
        "  def setup(self):\n",
        "    # Picking affordances.\n",
        "    self.pick_net = ResNet(1)\n",
        "\n",
        "    # Pick-conditioned placing affordances.\n",
        "    self.q_net = ResNet(3)  # Query (crop around pick location).\n",
        "    self.k_net = ResNet(3)  # Key (place features).\n",
        "    self.crop_size = 64\n",
        "    self.crop_conv = nn.Conv(features=1, kernel_size=(self.crop_size, self.crop_size), use_bias=False, dtype=jnp.float32, padding='SAME')\n",
        "\n",
        "  def __call__(self, x, text, p=None, train=True):\n",
        "    B, H, W, C = x.shape\n",
        "    pick_out = self.pick_net(x, text)  # (B, H, W, 1)\n",
        "\n",
        "    # Get key features.\n",
        "    k = self.k_net(x, text)\n",
        "\n",
        "    # Add 0-padding before cropping.\n",
        "    h = self.crop_size // 2\n",
        "    x_crop = jnp.pad(x, [(0, 0), (h, h), (h, h), (0, 0)], 'maximum')\n",
        "\n",
        "    # Get query features and convolve them over key features.\n",
        "    place_out = jnp.zeros((0, H, W, 1), jnp.float32)\n",
        "    for b in range(B):\n",
        "\n",
        "      # Get coordinates at center of crop.\n",
        "      if p is None:\n",
        "        pick_out_b = pick_out[b, ...]  # (H, W, 1)\n",
        "        pick_out_b = pick_out_b.flatten()  # (H * W,)\n",
        "        amax_i = jnp.argmax(pick_out_b)\n",
        "        v, u = jnp.unravel_index(amax_i, (H, W))\n",
        "      else:\n",
        "        v, u = p[b, :]\n",
        "\n",
        "      # Get query crop.\n",
        "      x_crop_b = jax.lax.dynamic_slice(x_crop, (b, v, u, 0), (1, self.crop_size, self.crop_size, x_crop.shape[3]))\n",
        "      # x_crop_b = x_crop[b:b+1, v:(v + self.crop_size), u:(u + self.crop_size), ...]\n",
        "\n",
        "      # Convolve q (query) across k (key).\n",
        "      q = self.q_net(x_crop_b, text[b:b+1, :])  # (1, H, W, 3)\n",
        "      q = jnp.transpose(q, (1, 2, 3, 0))  # (H, W, 3, 1)\n",
        "      place_out_b = self.crop_conv.apply({'params': {'kernel': q}}, k[b:b+1, ...])  # (1, H, W, 1)\n",
        "      scale = 1 / (self.crop_size * self.crop_size)  # For higher softmax temperatures.\n",
        "      place_out_b *= scale\n",
        "      place_out = jnp.concatenate((place_out, place_out_b), axis=0)\n",
        "\n",
        "    return pick_out, place_out\n",
        "\n",
        "\n",
        "def n_params(params):\n",
        "  return jnp.sum(jnp.int32([n_params(v) if isinstance(v, dict) or isinstance(v, flax.core.frozen_dict.FrozenDict) else np.prod(v.shape) for v in params.values()]))"
      ],
      "metadata": {
        "id": "gDE22S7ydjDN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Training**\n",
        "Train CLIP-conditioned Transporter Nets."
      ],
      "metadata": {
        "id": "4rqB5LrY3dXZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Compute CLIP features for text in the dataset.\n",
        "\n",
        "# Precompute CLIP features for all text in training dataset.\n",
        "text_tokens = clip.tokenize(dataset['text']).cuda()\n",
        "text_i = 0\n",
        "data_text_feats = np.zeros((0, 512), dtype=np.float32)\n",
        "while text_i < len(text_tokens):\n",
        "  batch_size = min(len(text_tokens) - text_i, 512)\n",
        "  text_batch = text_tokens[text_i:text_i+batch_size]\n",
        "  with torch.no_grad():\n",
        "    batch_feats = clip_model.encode_text(text_batch).float()\n",
        "  batch_feats /= batch_feats.norm(dim=-1, keepdim=True)\n",
        "  batch_feats = np.float32(batch_feats.cpu())\n",
        "  data_text_feats = np.concatenate((data_text_feats, batch_feats), axis=0)\n",
        "  text_i += batch_size"
      ],
      "metadata": {
        "id": "RgLWqlwqfY_x",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Define Transporter Nets train and eval functions\n",
        "\n",
        "# Train with InfoNCE loss over pick and place positions.\n",
        "@jax.jit\n",
        "def train_step(optimizer, batch):\n",
        "  def loss_fn(params):\n",
        "    batch_size = batch['img'].shape[0]\n",
        "    pick_logits, place_logits = TransporterNets().apply({'params': params}, batch['img'], batch['text'], batch['pick_yx'])\n",
        "\n",
        "    # InfoNCE pick loss.\n",
        "    pick_logits = pick_logits.reshape(batch_size, -1)\n",
        "    pick_onehot = batch['pick_onehot'].reshape(batch_size, -1)\n",
        "    pick_loss = jnp.mean(optax.softmax_cross_entropy(logits=pick_logits, labels=pick_onehot), axis=0)\n",
        "\n",
        "    # InfoNCE place loss.\n",
        "    place_logits = place_logits.reshape(batch_size, -1)\n",
        "    place_onehot = batch['place_onehot'].reshape(batch_size, -1)\n",
        "    place_loss = jnp.mean(optax.softmax_cross_entropy(logits=place_logits, labels=place_onehot), axis=0)\n",
        "    \n",
        "    loss = pick_loss + place_loss\n",
        "    return loss, (pick_logits, place_logits)\n",
        "  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n",
        "  (loss, logits), grad = grad_fn(optimizer.target)\n",
        "  optimizer = optimizer.apply_gradient(grad)\n",
        "  return optimizer, loss, grad, logits\n",
        "\n",
        "@jax.jit\n",
        "def eval_step(params, batch):\n",
        "  pick_logits, place_logits = TransporterNets().apply({'params': params}, batch['img'], batch['text'])\n",
        "  return pick_logits, place_logits\n",
        "\n",
        "# Coordinate map (i.e. position encoding).\n",
        "coord_x, coord_y = np.meshgrid(np.linspace(-1, 1, 224), np.linspace(-1, 1, 224), sparse=False, indexing='ij')\n",
        "coords = np.concatenate((coord_x[..., None], coord_y[..., None]), axis=2)"
      ],
      "metadata": {
        "id": "WUofYVFpjHnE",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown **TensorBoard:** Displays an interactive TensorBoard.\n",
        "name = datetime.datetime.now().strftime(f'%Y-%m-%d-%H:%M:%S-cliport')\n",
        "logdir = os.path.join(\"logs\", name)\n",
        "writer = tensorboard.SummaryWriter(logdir)\n",
        "%tensorboard --logdir logs"
      ],
      "metadata": {
        "cellView": "form",
        "id": "4tHOKjGM0_wh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Train your own model, or load a pretrained one.\n",
        "load_pretrained = True  #@param {type:\"boolean\"}\n",
        "\n",
        "# Initialize model weights using dummy tensors.\n",
        "rng = jax.random.PRNGKey(0)\n",
        "rng, key = jax.random.split(rng)\n",
        "init_img = jnp.ones((4, 224, 224, 5), jnp.float32)\n",
        "init_text = jnp.ones((4, 512), jnp.float32)\n",
        "init_pix = jnp.zeros((4, 2), np.int32)\n",
        "init_params = TransporterNets().init(key, init_img, init_text, init_pix)['params']\n",
        "print(f'Model parameters: {n_params(init_params):,}')\n",
        "optim = flax.optim.Adam(learning_rate=1e-4).create(init_params)\n",
        "\n",
        "if load_pretrained:\n",
        "  ckpt_path = f'ckpt_{40000}'\n",
        "  if not os.path.exists(ckpt_path):\n",
        "    !gdown --id 1Nq0q1KbqHOA5O7aRSu4u7-u27EMMXqgP\n",
        "  optim = checkpoints.restore_checkpoint(ckpt_path, optim)\n",
        "  print('Loaded:', ckpt_path)\n",
        "else:\n",
        "\n",
        "  # Training loop.\n",
        "  batch_size = 8\n",
        "  for train_iter in range(1, 40001):\n",
        "    batch_i = np.random.randint(dataset_size, size=batch_size)\n",
        "    text_feat = data_text_feats[batch_i, ...]\n",
        "    img = dataset['image'][batch_i, ...] / 255\n",
        "    img = np.concatenate((img, np.broadcast_to(coords[None, ...], (batch_size,) + coords.shape)), axis=3)\n",
        "\n",
        "    # Get onehot label maps.\n",
        "    pick_yx = np.zeros((batch_size, 2), dtype=np.int32)\n",
        "    pick_onehot = np.zeros((batch_size, 224, 224), dtype=np.float32)\n",
        "    place_onehot = np.zeros((batch_size, 224, 224), dtype=np.float32)\n",
        "    for i in range(len(batch_i)):\n",
        "      pick_y, pick_x  = dataset['pick_yx'][batch_i[i], :]\n",
        "      place_y, place_x = dataset['place_yx'][batch_i[i], :]\n",
        "      pick_onehot[i, pick_y, pick_x] = 1\n",
        "      place_onehot[i, place_y, place_x] = 1\n",
        "      # pick_onehot[i, ...] = scipy.ndimage.gaussian_filter(pick_onehot[i, ...], sigma=3)\n",
        "\n",
        "      # Data augmentation (random translation).\n",
        "      roll_y, roll_x = np.random.randint(-112, 112, size=2)\n",
        "      img[i, ...] = np.roll(img[i, ...], roll_y, axis=0)\n",
        "      img[i, ...] = np.roll(img[i, ...], roll_x, axis=1)\n",
        "      pick_onehot[i, ...] = np.roll(pick_onehot[i, ...], roll_y, axis=0)\n",
        "      pick_onehot[i, ...] = np.roll(pick_onehot[i, ...], roll_x, axis=1)\n",
        "      place_onehot[i, ...] = np.roll(place_onehot[i, ...], roll_y, axis=0)\n",
        "      place_onehot[i, ...] = np.roll(place_onehot[i, ...], roll_x, axis=1)\n",
        "      pick_yx[i, 0] = pick_y + roll_y\n",
        "      pick_yx[i, 1] = pick_x + roll_x\n",
        "\n",
        "    # Backpropagate.\n",
        "    batch = {}\n",
        "    batch['img'] = jnp.float32(img)\n",
        "    batch['text'] = jnp.float32(text_feat)\n",
        "    batch['pick_yx'] = jnp.int32(pick_yx)\n",
        "    batch['pick_onehot'] = jnp.float32(pick_onehot)\n",
        "    batch['place_onehot'] = jnp.float32(place_onehot)\n",
        "    rng, batch['rng'] = jax.random.split(rng)\n",
        "    optim, loss, _, _ = train_step(optim, batch)\n",
        "    writer.scalar('train/loss', loss, train_iter)\n",
        "\n",
        "    if train_iter % np.power(10, min(4, np.floor(np.log10(train_iter)))) == 0:\n",
        "      print(f'Train Step: {train_iter} Loss: {loss}')\n",
        "    \n",
        "    if train_iter % 1000 == 0:\n",
        "      checkpoints.save_checkpoint('.', optim, train_iter, prefix='ckpt_', keep=100000, overwrite=True)"
      ],
      "metadata": {
        "id": "vAGfC_yTd76V",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Demo:** CLIPort variant\n",
        "Test a variant of CLIPort: text-conditioned translation-only Transporter Nets.\n",
        "\n",
        "Text must generally be in the form: \"Pick the **{object}** and place it on the **{location}**.\"\n",
        "\n",
        "Admissible **objects:** \"{color} block\" (e.g. \"blue block\")\n",
        "\n",
        "Admissible **colors:** red, orange, yellow, green, blue, purple, pink, cyan, gray, brown\n",
        "\n",
        "Admissible **locations:** \"{color} block\" or \"{color} bowl\" or \"top/bottom/left/right side\" or \"top/bottom left/right corner\" or \"middle\""
      ],
      "metadata": {
        "id": "XOb79YDd3ijm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define and reset environment.\n",
        "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
        "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
        "\n",
        "np.random.seed(42)\n",
        "obs = env.reset(config)\n",
        "img = env.get_camera_image()\n",
        "plt.imshow(img)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "JBZORSgnywxn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "user_input = 'Pick the yellow block and place it on the blue bowl.'  #@param {type:\"string\"}\n",
        "\n",
        "# Show camera image before pick and place.\n",
        "\n",
        "def run_cliport(obs, text):\n",
        "  before = env.get_camera_image()\n",
        "  prev_obs = obs['image'].copy()\n",
        "\n",
        "  # Tokenize text and get CLIP features.\n",
        "  text_tokens = clip.tokenize(text).cuda()\n",
        "  with torch.no_grad():\n",
        "    text_feats = clip_model.encode_text(text_tokens).float()\n",
        "  text_feats /= text_feats.norm(dim=-1, keepdim=True)\n",
        "  text_feats = np.float32(text_feats.cpu())\n",
        "\n",
        "  # Normalize image and add batch dimension.\n",
        "  img = obs['image'][None, ...] / 255\n",
        "  img = np.concatenate((img, coords[None, ...]), axis=3)\n",
        "\n",
        "  # Run Transporter Nets to get pick and place heatmaps.\n",
        "  batch = {'img': jnp.float32(img), 'text': jnp.float32(text_feats)}\n",
        "  pick_map, place_map = eval_step(optim.target, batch)\n",
        "  pick_map, place_map = np.float32(pick_map), np.float32(place_map)\n",
        "\n",
        "  # Get pick position.\n",
        "  pick_max = np.argmax(np.float32(pick_map)).squeeze()\n",
        "  pick_yx = (pick_max // 224, pick_max % 224)\n",
        "  pick_yx = np.clip(pick_yx, 20, 204)\n",
        "  pick_xyz = obs['xyzmap'][pick_yx[0], pick_yx[1]]\n",
        "\n",
        "  # Get place position.\n",
        "  place_max = np.argmax(np.float32(place_map)).squeeze()\n",
        "  place_yx = (place_max // 224, place_max % 224)\n",
        "  place_yx = np.clip(place_yx, 20, 204)\n",
        "  place_xyz = obs['xyzmap'][place_yx[0], place_yx[1]]\n",
        "\n",
        "  # Step environment.\n",
        "  act = {'pick': pick_xyz, 'place': place_xyz}\n",
        "  obs, _, _, _ = env.step(act)\n",
        "\n",
        "  # Show pick and place action.\n",
        "  plt.title(text)\n",
        "  plt.imshow(prev_obs)\n",
        "  plt.arrow(pick_yx[1], pick_yx[0], place_yx[1]-pick_yx[1], place_yx[0]-pick_yx[0], color='w', head_starts_at_zero=False, head_width=7, length_includes_head=True)\n",
        "  plt.show()\n",
        "\n",
        "  # Show debug plots.\n",
        "  plt.subplot(1, 2, 1)\n",
        "  plt.title('Pick Heatmap')\n",
        "  plt.imshow(pick_map.reshape(224, 224))\n",
        "  plt.subplot(1, 2, 2)\n",
        "  plt.title('Place Heatmap')\n",
        "  plt.imshow(place_map.reshape(224, 224))\n",
        "  plt.show()\n",
        "\n",
        "  # Show video of environment rollout.\n",
        "  debug_clip = ImageSequenceClip(env.cache_video, fps=25)\n",
        "  display(debug_clip.ipython_display(autoplay=1, loop=1, center=False))\n",
        "  env.cache_video = []\n",
        "\n",
        "  # Show camera image after pick and place.\n",
        "  plt.subplot(1, 2, 1)\n",
        "  plt.title('Before')\n",
        "  plt.imshow(before)\n",
        "  plt.subplot(1, 2, 2)\n",
        "  plt.title('After')\n",
        "  after = env.get_camera_image()\n",
        "  plt.imshow(after)\n",
        "  plt.show()\n",
        "\n",
        "  # return pick_xyz, place_xyz, pick_map, place_map, pick_yx, place_yx\n",
        "  return obs\n",
        "\n",
        "\n",
        "# pick_xyz, place_xyz, pick_map, place_map, pick_yx, place_yx = \n",
        "obs = run_cliport(obs, user_input)\n"
      ],
      "metadata": {
        "id": "l0soZhUbyxUs",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Socratic Models:** **LM** (GPT-3) + **VLM** (ViLD) + **Policy** (CLIPort)"
      ],
      "metadata": {
        "id": "GO2qUyPblbdp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Define Socratic helper functions.\n",
        "\n",
        "gpt3_prompt = \"\"\"\n",
        "objects = [\"cyan block\", \"yellow block\", \"brown block\", \"green bowl\"]\n",
        "# move all the blocks to the top left corner.\n",
        "robot.pick_and_place(\"brown block\", \"top left corner\")\n",
        "robot.pick_and_place(\"cyan block\", \"top left corner\")\n",
        "robot.pick_and_place(\"yellow block\", \"top left corner\")\n",
        "# put the yellow one the green thing.\n",
        "robot.pick_and_place(\"yellow block\", \"green bowl\")\n",
        "# undo that.\n",
        "robot.pick_and_place(\"yellow block\", \"top left corner\")\n",
        "objects = [\"pink block\", \"gray block\", \"orange block\"]\n",
        "# move the pinkish colored block on the bottom side.\n",
        "robot.pick_and_place(\"pink block\", \"bottom side\")\n",
        "objects = [\"orange block\", \"purple bowl\", \"cyan block\", \"brown bowl\", \"pink block\"]\n",
        "# stack the blocks.\n",
        "robot.pick_and_place(\"pink block\", \"orange block\")\n",
        "robot.pick_and_place(\"cyan block\", \"pink block\")\n",
        "# unstack that.\n",
        "robot.pick_and_place(\"cyan block\", \"bottom left\")\n",
        "robot.pick_and_place(\"pink block\", \"left side\")\n",
        "objects = [\"red block\", \"brown block\", \"purple bowl\", \"gray bowl\", \"brown bowl\", \"pink block\", \"purple block\"]\n",
        "# group the brown objects together.\n",
        "robot.pick_and_place(\"brown block\", \"brown bowl\")\n",
        "objects = [\"orange bowl\", \"red block\", \"orange block\", \"red bowl\", \"purple bowl\", \"purple block\"]\n",
        "# sort all the blocks into their matching color bowls.\n",
        "robot.pick_and_place(\"orange block\", \"orange bowl\")\n",
        "robot.pick_and_place(\"red block\", \"red bowl\")\n",
        "robot.pick_and_place(\"purple block\", \"purple bowl\")\n",
        "\"\"\"\n",
        "\n",
        "gpt_version = \"text-davinci-002\"\n",
        "def LM(prompt, max_tokens=128, temperature=0, stop=None):\n",
        "  response = openai.Completion.create(engine=gpt_version, prompt=prompt, max_tokens=max_tokens, temperature=temperature, stop=stop)\n",
        "  return response[\"choices\"][0][\"text\"].strip()\n",
        "\n",
        "def VLM(img):\n",
        "  imageio.imwrite('tmp.jpg', img)\n",
        "  found_objects = vild(image_path, category_name_string, vild_params, plot_on=True)\n",
        "  return 'objects = [' + ', '.join([f'\\\"{o}\\\"' for o in found_objects]) + ']'"
      ],
      "metadata": {
        "cellView": "form",
        "id": "xYptWXd6J-V2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define and reset environment.\n",
        "config = {'pick':  ['yellow block', 'green block', 'blue block'],\n",
        "          'place': ['yellow bowl', 'green bowl', 'blue bowl']}\n",
        "np.random.seed(42)\n",
        "obs = env.reset(config)\n",
        "img = env.get_camera_image()\n",
        "plt.imshow(img)\n",
        "plt.show()\n",
        "\n",
        "# Reset context and get scene description from VLM.\n",
        "img = env.get_camera_image_top()\n",
        "img = np.flipud(img.transpose(1, 0, 2))\n",
        "scene_description = VLM(img)\n",
        "context = gpt3_prompt\n",
        "context += scene_description + '\\n'"
      ],
      "metadata": {
        "id": "uau2ee5ZJWd-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "user_input = 'sort the blocks in the way you think it best fits'  #@param {type:\"string\"}\n",
        "\n",
        "if user_input[-1] != '.':\n",
        "  user_input += '.'\n",
        "\n",
        "context += '# ' + user_input + '\\n'\n",
        "response = LM(context, stop=['#', 'objects ='])\n",
        "context += response + '\\n'\n",
        "\n",
        "step_text = ''\n",
        "def pick_and_place(obj1, obj2):\n",
        "  global step_text\n",
        "  step_text = f'Pick the {obj1} and place it on the {obj2}.'\n",
        "\n",
        "# Execute commands given by LM.\n",
        "step_cmds = response.split('\\n')\n",
        "print('LM generated plan:')\n",
        "for step_cmd in step_cmds:\n",
        "  step_cmd = step_cmd.replace('robot.', '')\n",
        "  exec(step_cmd)\n",
        "  print('Step:', step_text)\n",
        "  obs = run_cliport(obs, step_text)\n",
        "\n",
        "user_input = ''\n",
        "\n",
        "context += scene_description + '\\n'"
      ],
      "metadata": {
        "cellView": "form",
        "id": "eEuNP0_JKHnH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}